added groups

Author: ayasyrev

docs/Net.html

@@ -190,6 +190,60 @@ <h2 id="ResBlock" class="doc_header"><code>class</code> <code>ResBlock</code><a
 <div class="cell border-box-sizing code_cell rendered">
 <div class="input">
 
+<div class="inner_cell">
+    <div class="input_area">
+<div class=" highlight hl-ipython3"><pre><span></span><span class="n">ResBlock</span><span class="p">(</span><span class="mi">4</span><span class="p">,</span><span class="mi">64</span><span class="p">,</span><span class="mi">64</span><span class="p">,</span><span class="n">sa</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">groups</span><span class="o">=</span><span class="mi">4</span><span class="p">)</span>
+</pre></div>
+
+    </div>
+</div>
+</div>
+
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+
+<div class="output_text output_subarea output_execute_result">
+<pre>ResBlock(
+  (convs): Sequential(
+    (conv_0): ConvLayer(
+      (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
+      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+      (act_fn): ReLU(inplace=True)
+    )
+    (conv_1): ConvLayer(
+      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=4, bias=False)
+      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+      (act_fn): ReLU(inplace=True)
+    )
+    (conv_2): ConvLayer(
+      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
+      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+    )
+    (sa): SimpleSelfAttention(
+      (conv): Conv1d(256, 256, kernel_size=(1,), stride=(1,), bias=False)
+    )
+  )
+  (act_fn): ReLU(inplace=True)
+)</pre>
+</div>
+
+</div>
+
+</div>
+</div>
+
+</div>
+    {% endraw %}
+
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+<div class="input">
+
 <div class="inner_cell">
     <div class="input_area">
 <div class=" highlight hl-ipython3"><pre><span></span><span class="n">ResBlock</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span><span class="mi">64</span><span class="p">,</span><span class="mi">64</span><span class="p">,</span><span class="n">act_fn</span><span class="o">=</span><span class="n">nn</span><span class="o">.</span><span class="n">LeakyReLU</span><span class="p">(),</span> <span class="n">bn_1st</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
@@ -260,7 +314,7 @@ <h1 id="NewResBlock">NewResBlock<a class="anchor-link" href="#NewResBlock"> </a>
 
 
 <div class="output_markdown rendered_html output_subarea ">
-<h2 id="NewResBlock" class="doc_header"><code>class</code> <code>NewResBlock</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/net.py#L46" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>NewResBlock</code>(<strong><code>expansion</code></strong>, <strong><code>ni</code></strong>, <strong><code>nh</code></strong>, <strong><code>stride</code></strong>=<em><code>1</code></em>, <strong><code>conv_layer</code></strong>=<em><code>'ConvLayer'</code></em>, <strong><code>act_fn</code></strong>=<em><code>ReLU(inplace=True)</code></em>, <strong><code>zero_bn</code></strong>=<em><code>True</code></em>, <strong><code>bn_1st</code></strong>=<em><code>True</code></em>, <strong><code>pool</code></strong>=<em><code>AvgPool2d(kernel_size=2, stride=2, padding=0)</code></em>, <strong><code>sa</code></strong>=<em><code>False</code></em>, <strong><code>sym</code></strong>=<em><code>False</code></em>) :: <code>Module</code></p>
+<h2 id="NewResBlock" class="doc_header"><code>class</code> <code>NewResBlock</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/net.py#L46" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>NewResBlock</code>(<strong><code>expansion</code></strong>, <strong><code>ni</code></strong>, <strong><code>nh</code></strong>, <strong><code>stride</code></strong>=<em><code>1</code></em>, <strong><code>conv_layer</code></strong>=<em><code>'ConvLayer'</code></em>, <strong><code>act_fn</code></strong>=<em><code>ReLU(inplace=True)</code></em>, <strong><code>zero_bn</code></strong>=<em><code>True</code></em>, <strong><code>bn_1st</code></strong>=<em><code>True</code></em>, <strong><code>pool</code></strong>=<em><code>AvgPool2d(kernel_size=2, stride=2, padding=0)</code></em>, <strong><code>sa</code></strong>=<em><code>False</code></em>, <strong><code>sym</code></strong>=<em><code>False</code></em>, <strong><code>groups</code></strong>=<em><code>1</code></em>) :: <code>Module</code></p>
 </blockquote>
 <p>Base class for all neural network modules.</p>
 <p>Your models should also subclass this class.</p>
@@ -404,7 +458,10 @@ <h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="http
 
 
 <div class="output_text output_subarea output_execute_result">
-<pre> constr Net</pre>
+<pre> constr Net
+ expansion: 1, sa: 0, groups: 1
+ stem sizes: [3, 32, 32, 64]
+ body sizes [64, 64, 128, 256, 512]</pre>
 </div>
 
 </div>
@@ -461,7 +518,10 @@ <h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="http
     </div>
 </div>
 </div>
-
+<details class="description">
+      <summary data-open="Hide Output" data-close="Show Output"></summary>
+        <summary></summary>
+        
 <div class="output_wrapper">
 <div class="output">
 
@@ -473,18 +533,18 @@ <h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="http
 <pre>Sequential(
   (conv_0): ConvLayer(
     (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
+    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
     (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
-    (act_fn): ReLU(inplace=True)
   )
   (conv_1): ConvLayer(
     (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
     (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
-    (act_fn): ReLU(inplace=True)
   )
   (conv_2): ConvLayer(
     (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
     (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
-    (act_fn): ReLU(inplace=True)
   )
   (stem_pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
 )</pre>
@@ -495,6 +555,7 @@ <h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="http
 </div>
 </div>
 
+    </details>
 </div>
     {% endraw %}
 
@@ -562,7 +623,10 @@ <h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="http
     </div>
 </div>
 </div>
-
+<details class="description">
+      <summary data-open="Hide Output" data-close="Show Output"></summary>
+        <summary></summary>
+        
 <div class="output_wrapper">
 <div class="output">
 
@@ -613,6 +677,7 @@ <h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="http
 </div>
 </div>
 
+    </details>
 </div>
     {% endraw %}
 
@@ -647,73 +712,6 @@ <h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="http
 </div>
 </div>
 
-</div>
-    {% endraw %}
-
-    {% raw %}
-    
-<div class="cell border-box-sizing code_cell rendered">
-<div class="input">
-
-<div class="inner_cell">
-    <div class="input_area">
-<div class=" highlight hl-ipython3"><pre><span></span><span class="n">model</span><span class="o">.</span><span class="n">stem_bn_end</span> <span class="o">=</span> <span class="kc">True</span>
-</pre></div>
-
-    </div>
-</div>
-</div>
-
-</div>
-    {% endraw %}
-
-    {% raw %}
-    
-<div class="cell border-box-sizing code_cell rendered">
-<div class="input">
-
-<div class="inner_cell">
-    <div class="input_area">
-<div class=" highlight hl-ipython3"><pre><span></span><span class="n">model</span><span class="o">.</span><span class="n">stem</span>
-</pre></div>
-
-    </div>
-</div>
-</div>
-
-<div class="output_wrapper">
-<div class="output">
-
-<div class="output_area">
-
-
-
-<div class="output_text output_subarea output_execute_result">
-<pre>Sequential(
-  (conv_0): ConvLayer(
-    (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
-    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
-    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
-  )
-  (conv_1): ConvLayer(
-    (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
-    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
-    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
-  )
-  (conv_2): ConvLayer(
-    (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
-    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
-  )
-  (stem_pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
-  (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
-)</pre>
-</div>
-
-</div>
-
-</div>
-</div>
-
 </div>
     {% endraw %}
 
@@ -804,9 +802,6 @@ <h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="http
           (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
           (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
         )
-        (sa): SimpleSelfAttention(
-          (conv): Conv1d(256, 256, kernel_size=(1,), stride=(1,), bias=False)
-        )
       )
       (merge): LeakyReLU(negative_slope=0.01, inplace=True)
     )
@@ -955,23 +950,6 @@ <h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="http
 </div>
 
     </details>
-</div>
-    {% endraw %}
-
-    {% raw %}
-    
-<div class="cell border-box-sizing code_cell rendered">
-<div class="input">
-
-<div class="inner_cell">
-    <div class="input_area">
-<div class=" highlight hl-ipython3"><pre><span></span><span class="n">model</span><span class="o">.</span><span class="n">stem_sizes</span> <span class="o">=</span> <span class="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">32</span><span class="p">,</span> <span class="mi">32</span><span class="p">,</span> <span class="mi">64</span><span class="p">,</span> <span class="mi">64</span><span class="p">]</span>
-</pre></div>
-
-    </div>
-</div>
-</div>
-
 </div>
     {% endraw %}
 
@@ -1027,7 +1005,11 @@ <h2 id="xresnet-constructor">xresnet constructor<a class="anchor-link" href="#xr
 
 
 <div class="output_text output_subarea output_execute_result">
-<pre>( constr xresnet50, 10)</pre>
+<pre>( constr xresnet50
+  expansion: 4, sa: 0, groups: 1
+  stem sizes: [3, 32, 32, 64]
+  body sizes [16, 64, 128, 256, 512],
+ 10)</pre>
 </div>
 
 </div>

model_constructor/net.py

@@ -23,14 +23,15 @@ def init_cnn(m):
 class ResBlock(nn.Module):
     def __init__(self, expansion, ni, nh, stride=1,
                  conv_layer=ConvLayer, act_fn=act_fn, zero_bn=True, bn_1st=True,
-                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False,sym=False):
+                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False,sym=False, groups=1):
         super().__init__()
         nf,ni = nh*expansion,ni*expansion
         layers  = [(f"conv_0", conv_layer(ni, nh, 3, stride=stride, act_fn=act_fn, bn_1st=bn_1st)),
                    (f"conv_1", conv_layer(nh, nf, 3, zero_bn=zero_bn, act=False, bn_1st=bn_1st))
         ] if expansion == 1 else [
                    (f"conv_0",conv_layer(ni, nh, 1, act_fn=act_fn, bn_1st=bn_1st)),
-                   (f"conv_1",conv_layer(nh, nh, 3, stride=stride, act_fn=act_fn, bn_1st=bn_1st)),
+                   (f"conv_1",conv_layer(nh, nh, 3, stride=stride, act_fn=act_fn, bn_1st=bn_1st,
+                                         groups=int(nh/groups))),
                    (f"conv_2",conv_layer(nh, nf, 1, zero_bn=zero_bn, act=False, bn_1st=bn_1st))
         ]
         if sa: layers.append(('sa', SimpleSelfAttention(nf,ks=1,sym=sym)))
@@ -46,15 +47,15 @@ def forward(self, x): return self.act_fn(self.convs(x) + self.idconv(self.pool(x
 class NewResBlock(nn.Module):
     def __init__(self, expansion, ni, nh, stride=1,
                  conv_layer=ConvLayer, act_fn=act_fn, zero_bn=True, bn_1st=True,
-                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False,sym=False):
+                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False,sym=False, groups=1):
         super().__init__()
         nf,ni = nh*expansion,ni*expansion
         self.reduce = noop if stride==1 else pool
         layers  = [(f"conv_0", conv_layer(ni, nh, 3, stride=1, act_fn=act_fn, bn_1st=bn_1st)), # stride 1 !!!
                    (f"conv_1", conv_layer(nh, nf, 3, zero_bn=zero_bn, act=False, bn_1st=bn_1st))
         ] if expansion == 1 else [
                    (f"conv_0",conv_layer(ni, nh, 1, act_fn=act_fn, bn_1st=bn_1st)),
-                   (f"conv_1",conv_layer(nh, nh, 3, stride=1, act_fn=act_fn, bn_1st=bn_1st)), # stride 1 !!!
+                   (f"conv_1",conv_layer(nh, nh, 3, stride=1, act_fn=act_fn, bn_1st=bn_1st, groups=int(nh/groups))), # stride 1 !!!
                    (f"conv_2",conv_layer(nh, nf, 1, zero_bn=zero_bn, act=False, bn_1st=bn_1st))
         ]
         if sa: layers.append(('sa', SimpleSelfAttention(nf,ks=1,sym=sym)))
@@ -83,7 +84,7 @@ def _make_layer(self,expansion,ni,nf,blocks,stride,sa):
             [(f"bl_{i}", self.block(expansion, ni if i==0 else nf, nf,
                     stride if i==0 else 1, sa=sa if i==blocks-1 else False,
                     conv_layer=self.conv_layer, act_fn=self.act_fn, pool=self.pool,
-                                    zero_bn=self.zero_bn, bn_1st=self.bn_1st))
+                                    zero_bn=self.zero_bn, bn_1st=self.bn_1st, groups=self.groups))
               for i in range(blocks)]))
 
 # Cell
@@ -110,7 +111,7 @@ def __init__(self, expansion=1, layers=[2,2,2,2], c_in=3, c_out=1000, name='Net'
         self.name = name
         self.c_in, self.c_out,self.expansion,self.layers = c_in,c_out,expansion,layers # todo setter for expansion
         self.act_fn, self.pool, self.sa = act_fn, pool, sa
-
+        self.groups = 1
 
         self.stem_sizes = [c_in,32,32,64]
         self.stem_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
@@ -129,7 +130,7 @@ def __init__(self, expansion=1, layers=[2,2,2,2], c_in=3, c_out=1000, name='Net'
 
     @property
     def block_szs(self):
-        return [64//self.expansion,64,128,256,512] +[256]*(len(self.layers)-4)
+        return [self.stem_sizes[-1]//self.expansion,64,128,256,512] +[256]*(len(self.layers)-4)
 
     @property
     def stem(self):
@@ -154,7 +155,7 @@ def __call__(self):
         model.extra_repr = lambda : f"model {self.name}"
         return model
     def __repr__(self):
-        return f" constr {self.name}"
+        return f" constr {self.name}\n expansion: {self.expansion}, sa: {self.sa}, groups: {self.groups}\n stem sizes: {self.stem_sizes}\n body sizes {self.block_szs}"
 
 # Cell
 # me = sys.modules[__name__]

model_constructor/twist.py

@@ -22,7 +22,7 @@ class ConvTwist(nn.Module):
     groups_ch = 8
     def __init__(self, ni, nf,
                  ks=3, stride=1, padding=1, bias=False,
-                 groups=1, iters=1, init_max=0.7):
+                 groups=1, iters=1, init_max=0.7, **kvargs):
         super().__init__()
         self.same = ni==nf and stride==1
         self.groups = ni//self.groups_ch if self.use_groups else 1
@@ -108,7 +108,7 @@ class ConvLayerTwist(ConvLayer): # replace Conv2d by Twist
 class NewResBlockTwist(nn.Module):
     def __init__(self, expansion, ni, nh, stride=1,
                  conv_layer=ConvLayer, act_fn=act_fn, bn_1st=True,
-                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False, sym=False, zero_bn=True):
+                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False, sym=False, zero_bn=True, **kvargs):
         super().__init__()
         nf,ni = nh*expansion,ni*expansion
 #         conv_layer = ConvLayerTwist
@@ -134,7 +134,7 @@ def forward(self, x):
 class ResBlockTwist(nn.Module):
     def __init__(self, expansion, ni, nh, stride=1,
                  conv_layer=ConvLayer, act_fn=act_fn, zero_bn=True, bn_1st=True,
-                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False,sym=False):
+                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False,sym=False, **kvargs):
         super().__init__()
         nf,ni = nh*expansion,ni*expansion
 #         conv_layer = ConvLayerTwist