[Features] Support FP16 training

detrex/layers/multi_scale_deform_attn.py

@@ -324,9 +324,11 @@ def forward(
                     reference_points.shape[-1]
                 )
             )
+
+        # the original impl for fp32 training
         if torch.cuda.is_available() and value.is_cuda:
             output = MultiScaleDeformableAttnFunction.apply(
-                value,
+                value.to(torch.float32) if value.dtype==torch.float16 else value,
                 spatial_shapes,
                 level_start_index,
                 sampling_locations,
@@ -338,6 +340,9 @@ def forward(
                 value, spatial_shapes, sampling_locations, attention_weights
             )
 
+        if value.dtype==torch.float16:
+            output=output.to(torch.float16)
+
         output = self.output_proj(output)
 
         if not self.batch_first:

projects/dino/README.md

@@ -31,6 +31,14 @@ Hao Zhang, Feng Li, Shilong Liu, Lei Zhang, Hang Su, Jun Zhu, Lionel M. Ni, Heun
 <td align="center">100</td>
 <td align="center">49.2</td>
 <td align="center"> <a href="https://github.com/IDEA-Research/detrex-storage/releases/download/v0.2.0/dino_r50_4scale_12ep_49_2AP.pth">model</a></td>
+</tr>
+ <tr><td align="left">DINO-R50-4scale <b> with AMP</b></td>
+<td align="center">R-50</td>
+<td align="center">IN1k</td>
+<td align="center">12</td>
+<td align="center">100</td>
+<td align="center">49.1</td>
+<td align="center"> - </td>
 </tr>
 <!-- ROW: dino_r50_4scale_12ep -->
  <tr><td align="left"><a href="configs/dino_r50_5scale_12ep.py">DINO-R50-5scale</a></td>
@@ -251,6 +259,7 @@ Hao Zhang, Feng Li, Shilong Liu, Lei Zhang, Hang Su, Jun Zhu, Lionel M. Ni, Heun
 - `Swin-X-384` means the backbone pretrained resolution is `384 x 384` and `IN22k to In1k` means the model is pretrained on `ImageNet-22k` and finetuned on `ImageNet-1k`.
 - ViT backbone using MAE pretraining weights following [ViTDet](https://github.com/facebookresearch/detectron2/tree/main/projects/ViTDet)  which can be downloaded in [MAE](https://github.com/facebookresearch/mae). And it's not stable to train ViTDet-DINO without warmup lr-scheduler.
 - `Focal-LRF-3Level`: means using `Large-Receptive-Field (LRF)` and `Focal-Level` is setted to `3`, please refer to [FocalNet](https://github.com/microsoft/FocalNet) for more details about the backbone settings.
+- `with AMP`: means using mixed precision training.
 
 **Notable facts and caveats**: The position embedding of DINO in detrex is different from the original repo. We set the tempureture and offsets in `PositionEmbeddingSine` to `10000` and `-0.5` which may make the model converge a little bit faster in the early stage and get a slightly better results (about 0.1mAP) in 12 epochs settings.
 

tools/train_net.py

@@ -25,7 +25,6 @@
     SimpleTrainer,
     default_argument_parser,
     default_setup,
-    default_writers,
     hooks,
     launch,
 )
@@ -72,8 +71,8 @@ def __init__(
                 from torch.cuda.amp import GradScaler
 
                 grad_scaler = GradScaler()
-            self.grad_scaler = grad_scaler
-
+        self.grad_scaler = grad_scaler
+        
         # set True to use amp training
         self.amp = amp
 
@@ -98,8 +97,8 @@ def run_step(self):
         """
         If you want to do something with the losses, you can wrap the model.
         """
-        loss_dict = self.model(data)
         with autocast(enabled=self.amp):
+            loss_dict = self.model(data)
             if isinstance(loss_dict, torch.Tensor):
                 losses = loss_dict
                 loss_dict = {"total_loss": loss_dict}
@@ -135,6 +134,17 @@ def clip_grads(self, params):
                 **self.clip_grad_params,
             )
 
+    def state_dict(self):
+        ret = super().state_dict()
+        if self.grad_scaler and self.amp:
+            ret["grad_scaler"] = self.grad_scaler.state_dict()
+        return ret
+
+    def load_state_dict(self, state_dict):
+        super().load_state_dict(state_dict)
+        if self.grad_scaler and self.amp:
+            self.grad_scaler.load_state_dict(state_dict["grad_scaler"])
+
 
 def do_test(cfg, model):
     if "evaluator" in cfg.dataloader: