Commit end rush N

Author: arthurdouillard

inclearn/__init__.pyc

@@ -1 +1,4 @@
-from . import (cifar_resnet, densenet, my_resnet, my_resnet2, my_resnet_brn, resnet, ucir_resnet)
+from . import (
+    cifar_resnet, densenet, my_resnet, my_resnet2, my_resnet_brn, my_resnet_mcbn, my_resnet_mtl,
+    resnet, resnet_mtl, ucir_resnet, vgg
+)

inclearn/convnet/__init__.py

@@ -64,9 +64,10 @@ def __init__(self, inplanes, increase_dim=False, last_relu=False, downsampling="
         if increase_dim:
             if downsampling == "stride":
                 self.downsampler = DownsampleStride()
-                self.downsample = lambda x: self.pad(self.downsampler(x))
+                self._need_pad = True
             else:
-                self.downsample = DownsampleConv(inplanes, planes)
+                self.downsampler = DownsampleConv(inplanes, planes)
+                self._need_pad = False
 
         self.last_relu = last_relu
 
@@ -83,7 +84,9 @@ def forward(self, x):
         y = self.bn_b(y)
 
         if self.increase_dim:
-            x = self.downsample(x)
+            x = self.downsampler(x)
+            if self._need_pad:
+                x = self.pad(x)
 
         y = x + y
 
@@ -212,7 +215,7 @@ def __init__(
         )
 
         if pooling_config["type"] == "avg":
-            self.pool = nn.AvgPool2d(8)
+            self.pool = nn.AdaptiveAvgPool2d((1, 1))
         elif pooling_config["type"] == "weldon":
             self.pool = pooling.WeldonPool2d(**pooling_config)
         else:
@@ -222,22 +225,33 @@ def __init__(
         if final_layer in (True, "conv"):
             self.final_layer = nn.Conv2d(self.out_dim, self.out_dim, kernel_size=1, bias=False)
         elif isinstance(final_layer, dict):
-            if final_layer["type"] == "bn_relu_fc":
+            if final_layer["type"] == "one_layer":
                 self.final_layer = nn.Sequential(
-                    nn.BatchNorm1d(self.out_dim), nn.ReLU(),
+                    nn.BatchNorm1d(self.out_dim), nn.ReLU(inplace=True),
                     nn.Linear(self.out_dim, int(self.out_dim * final_layer["reduction_factor"]))
                 )
+                self.out_dim = int(self.out_dim * final_layer["reduction_factor"])
+            elif final_layer["type"] == "two_layers":
+                self.final_layer = nn.Sequential(
+                    nn.BatchNorm1d(self.out_dim), nn.ReLU(inplace=True),
+                    nn.Linear(self.out_dim, self.out_dim), nn.BatchNorm1d(self.out_dim),
+                    nn.ReLU(inplace=True),
+                    nn.Linear(self.out_dim, int(self.out_dim * final_layer["reduction_factor"]))
+                )
+                self.out_dim = int(self.out_dim * final_layer["reduction_factor"])
             else:
                 raise ValueError("Unknown final layer type {}.".format(final_layer["type"]))
         else:
-            self.final_layer = lambda x: x
+            self.final_layer = None
 
         for m in self.modules():
             if isinstance(m, nn.Conv2d):
                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
             elif isinstance(m, nn.BatchNorm2d):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
 
         if zero_residual:
             for m in self.modules():
@@ -289,7 +303,9 @@ def forward(self, x):
     def end_features(self, x):
         x = self.pool(x)
         x = x.view(x.size(0), -1)
-        x = self.final_layer(x)
+
+        if self.final_layer is not None:
+            x = self.final_layer(x)
 
         return x