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)
Support checkpoint in DINO
Environment:
A100-40GB
,batch-size=1
,Swin-Large-384-4Scale-DINO
Usage
Swin-DINO
:Simple tutorial on adding checkpoint
We use
TransformerLayerSequence
andBaseTransformerLayer
as the basic block for DETR models in detrex, all you need to do is to wrap the layer inTransformerLayerSequence
withcheckpoint_wrapper
implemented in fairscale: