[Enhance]: Avoid overflow issue in pooling layers of RTMDet when usin…

…g AMP (open-mmlab#9670) Co-authored-by: Wenwei Zhang <[email protected]>
imclab · Jan 30, 2023 · 4e83e86 · 4e83e86
1 parent a2b3e84
commit 4e83e86
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/mmdet/models/backbones/csp_darknet.py b/mmdet/models/backbones/csp_darknet.py
@@ -115,7 +115,9 @@ def __init__(self,
 
     def forward(self, x):
         x = self.conv1(x)
-        x = torch.cat([x] + [pooling(x) for pooling in self.poolings], dim=1)
+        with torch.cuda.amp.autocast(enabled=False):
+            x = torch.cat(
+                [x] + [pooling(x) for pooling in self.poolings], dim=1)
         x = self.conv2(x)
         return x
 

diff --git a/mmdet/models/layers/se_layer.py b/mmdet/models/layers/se_layer.py
@@ -152,7 +152,8 @@ def __init__(self, channels: int, init_cfg: OptMultiConfig = None) -> None:
 
     def forward(self, x: Tensor) -> Tensor:
         """Forward function for ChannelAttention."""
-        out = self.global_avgpool(x)
+        with torch.cuda.amp.autocast(enabled=False):
+            out = self.global_avgpool(x)
         out = self.fc(out)
         out = self.act(out)
         return x * out