stem pool fix

Author: ayasyrev

model_constructor/model_constructor.py

@@ -66,7 +66,8 @@ def _make_stem(self):
                                           bn_layer=(not self.stem_bn_end) if i == (len(self.stem_sizes) - 2) else True,
                                           act_fn=self.act_fn, bn_1st=self.bn_1st))
             for i in range(len(self.stem_sizes) - 1)]
-    stem.append(('stem_pool', self.stem_pool))
+    if self.stem_pool is not None:
+        stem.append(('stem_pool', self.stem_pool))
     if self.stem_bn_end:
         stem.append(('norm', self.norm(self.stem_sizes[-1])))
     return nn.Sequential(OrderedDict(stem))
@@ -84,9 +85,10 @@ def _make_layer(self, expansion, ni, nf, blocks, stride, sa):
 
 
 def _make_body(self):
+    stride = 2 if self.stem_pool is None else 1  # if no pool on stem - stride = 2 for first block in body
     blocks = [(f"l_{i}", self._make_layer(self, self.expansion,
                                           ni=self.block_sizes[i], nf=self.block_sizes[i + 1],
-                                          blocks=l, stride=1 if i == 0 else 2,
+                                          blocks=l, stride=stride if i == 0 else 2,
                                           sa=self.sa if i == 0 else False))
               for i, l in enumerate(self.layers)]
     return nn.Sequential(OrderedDict(blocks))
@@ -114,7 +116,7 @@ def __init__(self, name='MC', c_in=3, c_out=1000,
                  zero_bn=True,
                  stem_stride_on=0,
                  stem_sizes=[32, 32, 64],
-                 stem_pool=nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
+                 stem_pool=nn.MaxPool2d(kernel_size=3, stride=2, padding=1),  # if stem_pool is None - no pool at stem
                  stem_bn_end=False,
                  _init_cnn=init_cnn,
                  _make_stem=_make_stem,
@@ -128,6 +130,8 @@ def __init__(self, name='MC', c_in=3, c_out=1000,
         del params['self']
         self.__dict__ = params
         self._block_sizes = params['block_sizes']
+        if type(self.stem_pool) is str:  # Hydra pass string value
+            self.stem_pool = None
         if self.stem_sizes[0] != self.c_in:
             self.stem_sizes = [self.c_in] + self.stem_sizes