Match BVLC GoogLeNet zero initialization of classifier

pytorch · Jan 9, 2019 · 8bc54e2 · 8bc54e2
1 parent 2c8caab
commit 8bc54e2
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/torchvision/models/googlenet.py b/torchvision/models/googlenet.py
@@ -79,6 +79,13 @@ def __init__(self, num_classes=1000, aux_logits=True, batch_norm=False):
                 if m.bias is not None:
                     nn.init.constant_(m.bias, 0.2)
 
+        # zero init classifier
+        for m in self.modules():
+            if isinstance(m, InceptionAux):
+                nn.init.zeros_(m.fc2.bias)
+            elif m == self.fc:
+                nn.init.zeros_(m.bias)
+
     def forward(self, x):
         x = self.conv1(x)
         x = self.maxpool1(x)