Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Features] Support encoder-decoder checkpoint in DINO #200

Merged
merged 2 commits into from
Feb 4, 2023

Conversation

rentainhe
Copy link
Collaborator

@rentainhe rentainhe commented Feb 4, 2023

Support checkpoint in DINO

Environment: A100-40GB, batch-size=1, Swin-Large-384-4Scale-DINO

  • without checkpoint: 20GB~22GB
  • with backbone (Swin-Large) checkpoint: 13GB~14GB
  • with encoder-decoder checkpoint: 16GB ~ 18GB
  • with backbone checkpoint + encoder-decoder checkpoint: 10~11GB

Usage

  • Open checkpoint in command line for Swin-DINO:
python tools/train_net.py \
    --config-file projects/dino/configs/dino_swin_large_384_4scale_12ep.py \
    --num-gpus 8 \
    train.init_checkpoint="path/to/checkpoint" \
    model.backbone.use_checkpoint=True \
    model.transformer.encoder.use_checkpoint=True \
    model.transformer.decoder.use_checkpoint=True \

Simple tutorial on adding checkpoint

We use TransformerLayerSequence and BaseTransformerLayer as the basic block for DETR models in detrex, all you need to do is to wrap the layer in TransformerLayerSequence with checkpoint_wrapper implemented in fairscale:

from fairscale.nn.checkpoint import checkpoint_wrapper

class DINOTransformerEncoder(TransformerLayerSequence):
    def __init__(*args, **kwargs):
        ...
        if use_checkpoint:
            for layer in self.layers:
                layer = checkpoint_wrapper(layer)

@zengzhaoyang zengzhaoyang merged commit dd5874d into main Feb 4, 2023
@rentainhe rentainhe deleted the support_dino_checkpoint branch February 4, 2023 05:50
Lontoone pushed a commit to Lontoone/detrex that referenced this pull request Jan 8, 2024
)

* support encoder decoder checkpoint in dino

* refine config

---------

Co-authored-by: ntianhe ren <rentianhe@dgx061.scc.idea>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants