fix bn error

Author: ayasyrev

README.md

@@ -3,7 +3,7 @@
 #################################################
 ### THIS FILE WAS AUTOGENERATED! DO NOT EDIT! ###
 #################################################
-# file to edit: index.ipynb
+# file to edit: nbs/index.ipynb
 # command to build the docs after a change: nbdev_build_docs
 
 -->

docs/constructor.html

@@ -12,7 +12,7 @@
 #################################################
 ### THIS FILE WAS AUTOGENERATED! DO NOT EDIT! ###
 #################################################
-# file to edit: 00_constructor.ipynb
+# file to edit: nbs/00_constructor.ipynb
 # command to build the docs after a change: nbdev_build_docs
 
 -->
@@ -80,8 +80,8 @@ <h2 id="Stem" class="doc_header"><code>class</code> <code>Stem</code><a href="ht
   sizes: [3, 64]
   (conv0): ConvLayer(
     (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
-    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
     (act_fn): ReLU(inplace=True)
+    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
   )
   (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
 )</pre>
@@ -279,6 +279,49 @@ <h2 id="Stem" class="doc_header"><code>class</code> <code>Stem</code><a href="ht
 </div>
 </div>
 
+</div>
+<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">stem</span> <span class="o">=</span> <span class="n">Stem</span><span class="p">(</span><span class="n">use_bn</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">stem_sizes</span><span class="o">=</span><span class="p">[</span><span class="mi">32</span><span class="p">],</span> <span class="n">stride_on</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">bn_1st</span><span class="o">=</span><span class="kc">False</span><span class="p">)</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>Stem(
+  sizes: [3, 32, 64]
+  (conv0): ConvLayer(
+    (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
+    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+    (act_fn): ReLU(inplace=True)
+  )
+  (conv1): ConvLayer(
+    (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
+    (act_fn): ReLU(inplace=True)
+  )
+  (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
+  (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
+)</pre>
+</div>
+
+</div>
+
+</div>
+</div>
+
 </div>
 <div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
 <div class="text_cell_render border-box-sizing rendered_html">
@@ -304,7 +347,7 @@ <h2 id="BasicBlock">BasicBlock<a class="anchor-link" href="#BasicBlock">&#182;</
 
 
 <div class="output_markdown rendered_html output_subarea ">
-<h2 id="BasicBlock" class="doc_header"><code>class</code> <code>BasicBlock</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L30" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>BasicBlock</code>(<strong><code>ni</code></strong>, <strong><code>nf</code></strong>, <strong><code>expansion</code></strong>=<em><code>1</code></em>, <strong><code>stride</code></strong>=<em><code>1</code></em>, <strong><code>bn_1st</code></strong>=<em><code>False</code></em>, <strong><code>zero_bn</code></strong>=<em><code>False</code></em>, <strong><code>conv_layer</code></strong>=<em><code>'ConvLayer'</code></em>, <strong>**<code>kwargs</code></strong>) :: <code>Module</code></p>
+<h2 id="BasicBlock" class="doc_header"><code>class</code> <code>BasicBlock</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L31" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>BasicBlock</code>(<strong><code>ni</code></strong>, <strong><code>nf</code></strong>, <strong><code>expansion</code></strong>=<em><code>1</code></em>, <strong><code>stride</code></strong>=<em><code>1</code></em>, <strong><code>bn_1st</code></strong>=<em><code>False</code></em>, <strong><code>zero_bn</code></strong>=<em><code>False</code></em>, <strong><code>conv_layer</code></strong>=<em><code>'ConvLayer'</code></em>, <strong>**<code>kwargs</code></strong>) :: <code>Module</code></p>
 </blockquote>
 <p>Basic block (simplified) as in pytorch resnet</p>
 
@@ -526,34 +569,6 @@ <h2 id="BasicBlock" class="doc_header"><code>class</code> <code>BasicBlock</code
 <h2 id="Body-constructor">Body constructor<a class="anchor-link" href="#Body-constructor">&#182;</a></h2>
 </div>
 </div>
-</div>
-<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="c1"># class Body(nn.Sequential): </span>
-<span class="c1">#     &#39;&#39;&#39;Constructor for body&#39;&#39;&#39;</span>
-<span class="c1">#     def __init__(self, block, </span>
-<span class="c1">#                  body_in=64, body_out=512, </span>
-<span class="c1">#                  layer_szs=[64,128,256,], blocks=[2,2,2,2],</span>
-<span class="c1">#                  expansion=1, **kwargs):  </span>
-<span class="c1">#         layer_szs = [body_in//expansion] + layer_szs + [body_out]</span>
-<span class="c1">#         num_layers = len(layer_szs)-1</span>
-<span class="c1">#         layers = [(f&quot;layer_{i}&quot;, self._make_layer(block, layer_szs[i], layer_szs[i+1], blocks[i], 1 if i==0 else 2, **kwargs))</span>
-<span class="c1">#                     for i in range(num_layers)]</span>
-<span class="c1">#         super().__init__(OrderedDict(layers))</span>
-<span class="c1">#     def _make_layer(self, block, ni, nf, blocks, stride, **kwargs):</span>
-<span class="c1">#         return nn.Sequential(OrderedDict(</span>
-<span class="c1">#             [(f&#39;block_{i}&#39;, block(ni if i==0 else nf, nf, stride if i==0 else 1, **kwargs))</span>
-<span class="c1">#               for i in range(blocks)]))</span>
-                    
-</pre></div>
-
-    </div>
-</div>
-</div>
-
 </div>
 <div class="cell border-box-sizing code_cell rendered">
 
@@ -567,7 +582,7 @@ <h2 id="Body-constructor">Body constructor<a class="anchor-link" href="#Body-con
 
 
 <div class="output_markdown rendered_html output_subarea ">
-<h2 id="BasicLayer" class="doc_header"><code>class</code> <code>BasicLayer</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L54" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>BasicLayer</code>(<strong><code>block</code></strong>, <strong><code>blocks</code></strong>, <strong><code>ni</code></strong>, <strong><code>nf</code></strong>, <strong><code>expansion</code></strong>, <strong><code>stride</code></strong>, <strong>**<code>kwargs</code></strong>) :: <code>Sequential</code></p>
+<h2 id="BasicLayer" class="doc_header"><code>class</code> <code>BasicLayer</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L55" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>BasicLayer</code>(<strong><code>block</code></strong>, <strong><code>blocks</code></strong>, <strong><code>ni</code></strong>, <strong><code>nf</code></strong>, <strong><code>expansion</code></strong>, <strong><code>stride</code></strong>, <strong>**<code>kwargs</code></strong>) :: <code>Sequential</code></p>
 </blockquote>
 <p>Layer from blocks</p>
 
@@ -588,7 +603,7 @@ <h2 id="BasicLayer" class="doc_header"><code>class</code> <code>BasicLayer</code
 
 
 <div class="output_markdown rendered_html output_subarea ">
-<h2 id="Body" class="doc_header"><code>class</code> <code>Body</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L68" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>Body</code>(<strong><code>block</code></strong>, <strong><code>body_in</code></strong>=<em><code>64</code></em>, <strong><code>body_out</code></strong>=<em><code>512</code></em>, <strong><code>bodylayer</code></strong>=<em><code>'BasicLayer'</code></em>, <strong><code>expansion</code></strong>=<em><code>1</code></em>, <strong><code>layer_szs</code></strong>=<em><code>[64, 128, 256]</code></em>, <strong><code>blocks</code></strong>=<em><code>[2, 2, 2, 2]</code></em>, <strong>**<code>kwargs</code></strong>) :: <code>Sequential</code></p>
+<h2 id="Body" class="doc_header"><code>class</code> <code>Body</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L69" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>Body</code>(<strong><code>block</code></strong>, <strong><code>body_in</code></strong>=<em><code>64</code></em>, <strong><code>body_out</code></strong>=<em><code>512</code></em>, <strong><code>bodylayer</code></strong>=<em><code>'BasicLayer'</code></em>, <strong><code>expansion</code></strong>=<em><code>1</code></em>, <strong><code>layer_szs</code></strong>=<em><code>[64, 128, 256]</code></em>, <strong><code>blocks</code></strong>=<em><code>[2, 2, 2, 2]</code></em>, <strong>**<code>kwargs</code></strong>) :: <code>Sequential</code></p>
 </blockquote>
 <p>Constructor for body</p>
 
@@ -846,7 +861,7 @@ <h1 id="Head">Head<a class="anchor-link" href="#Head">&#182;</a></h1>
 
 
 <div class="output_markdown rendered_html output_subarea ">
-<h2 id="Head" class="doc_header"><code>class</code> <code>Head</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L85" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>Head</code>(<strong><code>ni</code></strong>, <strong><code>nf</code></strong>, <strong>**<code>kwargs</code></strong>) :: <code>Sequential</code></p>
+<h2 id="Head" class="doc_header"><code>class</code> <code>Head</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L86" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>Head</code>(<strong><code>ni</code></strong>, <strong><code>nf</code></strong>, <strong>**<code>kwargs</code></strong>) :: <code>Sequential</code></p>
 </blockquote>
 <p>base head</p>
 
@@ -920,7 +935,7 @@ <h1 id="class-Net">class Net<a class="anchor-link" href="#class-Net">&#182;</a><
 
 
 <div class="output_markdown rendered_html output_subarea ">
-<h4 id="init_model" class="doc_header"><code>init_model</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L95" class="source_link" style="float:right">[source]</a></h4><blockquote><p><code>init_model</code>(<strong><code>model</code></strong>, <strong><code>zero_bn</code></strong>=<em><code>False</code></em>)</p>
+<h4 id="init_model" class="doc_header"><code>init_model</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L96" class="source_link" style="float:right">[source]</a></h4><blockquote><p><code>init_model</code>(<strong><code>model</code></strong>, <strong><code>zero_bn</code></strong>=<em><code>False</code></em>)</p>
 </blockquote>
 <p>Init model</p>
 
@@ -944,7 +959,7 @@ <h4 id="init_model" class="doc_header"><code>init_model</code><a href="https://g
 
 
 <div class="output_markdown rendered_html output_subarea ">
-<h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L102" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>Net</code>(<strong><code>stem</code></strong>=<em><code>'Stem'</code></em>, <strong><code>body</code></strong>=<em><code>'Body'</code></em>, <strong><code>block</code></strong>=<em><code>'BasicBlock'</code></em>, <strong><code>layer_szs</code></strong>=<em><code>[64, 128, 256]</code></em>, <strong><code>blocks</code></strong>=<em><code>[2, 2, 2, 2]</code></em>, <strong><code>head</code></strong>=<em><code>'Head'</code></em>, <strong><code>c_in</code></strong>=<em><code>3</code></em>, <strong><code>num_classes</code></strong>=<em><code>1000</code></em>, <strong><code>body_in</code></strong>=<em><code>64</code></em>, <strong><code>body_out</code></strong>=<em><code>512</code></em>, <strong><code>expansion</code></strong>=<em><code>1</code></em>, <strong><code>bn_1st</code></strong>=<em><code>False</code></em>, <strong><code>init_type</code></strong>=<em><code>'normal'</code></em>, <strong>**<code>kwargs</code></strong>) :: <code>Sequential</code></p>
+<h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/constructor.py#L103" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>Net</code>(<strong><code>stem</code></strong>=<em><code>'Stem'</code></em>, <strong><code>body</code></strong>=<em><code>'Body'</code></em>, <strong><code>block</code></strong>=<em><code>'BasicBlock'</code></em>, <strong><code>layer_szs</code></strong>=<em><code>[64, 128, 256]</code></em>, <strong><code>blocks</code></strong>=<em><code>[2, 2, 2, 2]</code></em>, <strong><code>head</code></strong>=<em><code>'Head'</code></em>, <strong><code>c_in</code></strong>=<em><code>3</code></em>, <strong><code>num_classes</code></strong>=<em><code>1000</code></em>, <strong><code>body_in</code></strong>=<em><code>64</code></em>, <strong><code>body_out</code></strong>=<em><code>512</code></em>, <strong><code>expansion</code></strong>=<em><code>1</code></em>, <strong><code>bn_1st</code></strong>=<em><code>False</code></em>, <strong><code>init_type</code></strong>=<em><code>'normal'</code></em>, <strong>**<code>kwargs</code></strong>) :: <code>Sequential</code></p>
 </blockquote>
 <p>Constructor for model</p>
 
@@ -1182,6 +1197,22 @@ <h2 id="Net" class="doc_header"><code>class</code> <code>Net</code><a href="http
 </div>
 </div>
 
+</div>
+<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">bs_test</span> <span class="o">=</span> <span class="mi">16</span>
+<span class="n">xb</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="n">bs_test</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">)</span>
+<span class="n">y</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="n">xb</span><span class="p">)</span>
+<span class="n">y</span><span class="o">.</span><span class="n">shape</span>
+</pre></div>
+
+    </div>
+</div>
+</div>
+
 <div class="output_wrapper">
 <div class="output">
 

docs/index.html

@@ -12,7 +12,7 @@
 #################################################
 ### THIS FILE WAS AUTOGENERATED! DO NOT EDIT! ###
 #################################################
-# file to edit: index.ipynb
+# file to edit: nbs/index.ipynb
 # command to build the docs after a change: nbdev_build_docs
 
 -->

docs/layers.html

@@ -12,7 +12,7 @@
 #################################################
 ### THIS FILE WAS AUTOGENERATED! DO NOT EDIT! ###
 #################################################
-# file to edit: 01_layers.ipynb
+# file to edit: nbs/01_layers.ipynb
 # command to build the docs after a change: nbdev_build_docs
 
 -->

model_constructor/constructor.py

@@ -18,8 +18,9 @@ def __init__(self, c_in=3,  stem_sizes=[], stem_out=64,
         num_layers = len(self.sizes)-1
         stem = [(f"conv{i}", conv_layer(self.sizes[i], self.sizes[i+1],
                 stride=2 if i==stride_on else 1, act=True,
-                bn_layer=not use_bn if i==num_layers-1 else True, **kwargs
-                )) for i in range(num_layers)]
+                bn_layer=not use_bn if i==num_layers-1 else True,
+                bn_1st=bn_1st, **kwargs))
+                    for i in range(num_layers)]
         if pool: stem += [('pool', nn.MaxPool2d(kernel_size=3, stride=2, padding=1))]
         if use_bn: stem.append(('bn', nn.BatchNorm2d(stem_out)))
         super().__init__(OrderedDict(stem))
@@ -109,8 +110,6 @@ def __init__(self,  stem=Stem,
                  body_in=64, body_out=512, expansion=1,
                  bn_1st=False,
                 init_type='normal', **kwargs):
-        # c_in = 3
-        # block_szs  = [64,128,128,256,256,512]
         super().__init__(OrderedDict([
             ('stem', stem(c_in=c_in,stem_out=body_in, **kwargs)),
             ('body', body(block, body_in, body_out,