[Features] Support encoder-decoder checkpoint in DINO

[rentainhe]

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)