Refactored DETR implementation on `optax` API.

google-research / scenic

Scenic: A Jax Library for Computer Vision Research and Beyond

Apache License 2.0

3.28k stars 428 forks source link

Refactored DETR implementation on `optax` API. #1062

Open MasterSkepticista opened 4 months ago

MasterSkepticista commented 4 months ago

Many functions and calls within existing detr implementation are deprecated in Flax as well as OTT.

This PR includes two major changes (all changes only under projects/baselines/detr):

Migrates from flax.optim to optax.
sinkhorn solver now uses updated ott-jax calls.

I have tested both configs for train/eval/checkpointing.

durianer-D commented 1 month ago

Hi @MasterSkepticista, I am experimenting with your modified detr, I have encountered many environment conflicts and tried to resolve them, but there are still some difficult bugs, can you provide your environment as a reference. Thank you very much.

MasterSkepticista commented 1 month ago

@durianer-D I created a new repo with a minimal implementation. It is also significantly faster to train than this PR. https://github.com/MasterSkepticista/detr.