fix bug caused by the comment duplicated BasicConv2d

alibaba · Jul 9, 2024 · 83ea6b7 · 83ea6b7
1 parent 1cf7cfe
commit 83ea6b7
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/easycv/models/backbones/inceptionv4.py b/easycv/models/backbones/inceptionv4.py
@@ -15,7 +15,12 @@
 
 class BasicConv2d(nn.Module):
 
-    def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0):
+    def __init__(self,
+                 in_planes,
+                 out_planes,
+                 kernel_size,
+                 stride=1,
+                 padding=0):
         super(BasicConv2d, self).__init__()
         self.conv = nn.Conv2d(
             in_planes,
@@ -357,9 +362,8 @@ def init_weights(self):
                 nn.init.constant_(m.bias, 0)
 
     def logits(self, features):
-        # Allows image of any size to be processed
-        adaptiveAvgPoolWidth = features.shape[2]
-        x = F.avg_pool2d(features, kernel_size=adaptiveAvgPoolWidth)
+        x = F.adaptive_avg_pool2d(features, output_size=(1, 1))
+        # x = F.avg_pool2d(features, kernel_size=adaptiveAvgPoolWidth)
         x = x.view(x.size(0), -1)  # B x 1536
         x = self.fc(x)
         # B x num_classes