Creating skip connections across nn.Module's

username: neel_g
cooked: <p>Hi, <img src="https://discuss.pytorch.org/images/emoji/apple/wave.png?v=12" title=":wave:" class="emoji" alt=":wave:" loading="lazy" width="20" height="20"> I was looking to implement U-Net Like long skip-connections across the encoder as well as the decoder. This <a href="https://github.com/rosinality/vq-vae-2-pytorch/blob/ef5f67c46f93624163776caec9e0d95063910eca/vqvae.py#L98-L158" rel="noopener nofollow ugc">snippet</a> illustrates my problem.</p>
<pre data-code-wrap="py"><code class="lang-py">blocks.extend( [nn.BatchNorm2d(self.channel), nn.ReLU(inplace=True)] ) #adding BatchNorm2d

self.blocks = nn.Sequential(*blocks)

def forward(self, input):
    return self.blocks(input)
</code></pre>
<p>Because I have 50+ layers in each <code>Module</code>, I can’t expose them individually. But otherwise, I don’t see how I can add skip connections <strong>across</strong> the encoder to the decoder, like <code>U-NET</code></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/7/27473a86c609b51291cfb99f59efc6dea81b0a34.png" data-download-href="https://discuss.pytorch.org/uploads/default/27473a86c609b51291cfb99f59efc6dea81b0a34" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/27473a86c609b51291cfb99f59efc6dea81b0a34_2_690x450.png" alt="image" data-base62-sha1="5Bte8hdAKZTRhQKBToorpIMhFVa" width="690" height="450" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/27473a86c609b51291cfb99f59efc6dea81b0a34_2_690x450.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/27473a86c609b51291cfb99f59efc6dea81b0a34_2_1035x675.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/27473a86c609b51291cfb99f59efc6dea81b0a34_2_1380x900.png 2x" data-dominant-color="F5F7F9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1412×922 102 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Does anyone have any ideas how I can simply do it? I suppose something like a skip connection from <code>i</code>th layer in encoder to <code>len(decoder)-i</code>th layer in the decoder (or maybe a bit complex like U-net, you get the idea)</p>
<p>P.S:- I am using a VQ-VAE-2, but for all intents and purposes its mostly a U-Net added with a quantization layer - so there shouldn’t be any major hiccups <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>
<p>Thanks! <img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=12" title=":hugs:" class="emoji" alt=":hugs:" loading="lazy" width="20" height="20"></p>

username: ptrblck
cooked: <p>The common approach would be to store the output activations of the encoder layers and reuse them in the decoder path directly. If you are seeing issues using this approach due to the large number of layers, I would suggest to store the outputs in e.g. a <code>list</code> and use a <code>loop</code> to execute the blocks of both parts of the model.<br>
E.g. something like this might work:</p>
<pre><code class="lang-python">def forward(self, x):
    out = x
    acts = []
    for layer in self.encoder:
        out = layer(out)
        acts.append(out)

    for a, layer in zip(acts[::-1], self.decoder):
        out = layer(out, a)

    return out
</code></pre>

username: neel_g
cooked: <p>thanks! even if I selectively append layer activations in the list, won’t it still prove a bottleneck for each forward pass as I’ve to doubly iterate, append and compute? <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"></p>

username: ptrblck
cooked: <p>I’m not sure why the for loop would create a bottleneck compared to the sequential execution of the layers. Could you explain your concern a bit more, please?</p>

username: neel_g
cooked: <p>I suppose I was concerned with having so many layers being appended to a <code>list</code> and iterated over - recomputing the same list every single time for each forward pass. I suppose it isn’t really that big of a deal-breaker since I am in the fine-tuning stage, but perhaps someone else may not find the solution to be optimal for their high-performance needs <img src="https://discuss.pytorch.org/images/emoji/apple/man_shrugging.png?v=12" title=":man_shrugging:" class="emoji" alt=":man_shrugging:" loading="lazy" width="20" height="20"></p>

username: ptrblck
cooked: <p>If I understand your concern correctly you re afraid that e.g.</p>
<pre><code class="lang-python">for module in modules:
    out = module(out)
</code></pre>
<p>would be slower than:</p>
<pre><code class="lang-python">out = module1(out)
out = module2(out)
out = module3(out)
...
out = module50(out)
</code></pre>
<p>If so, then I would claim that the Python overhead (if any) should be small compared to the actual workload execution, but you should certainly profile it if this use case fits your concern.</p>

username: neel_g
cooked: <p>lovely! <img src="https://discuss.pytorch.org/images/emoji/apple/ok_hand.png?v=12" title=":ok_hand:" class="emoji" alt=":ok_hand:" loading="lazy" width="20" height="20"> <img src="https://discuss.pytorch.org/images/emoji/apple/100.png?v=12" title=":100:" class="emoji" alt=":100:" loading="lazy" width="20" height="20"><br>
Just one last question - for your snippet above, we create a skip connection from <em>every</em> encoder layer to <em>every</em> decoder layer. My network is mostly composed of <code>resblocks</code> on both the side (and the task is <code>semantic segmentation</code>)</p>
<p>In that case, should I just add a few (say 7 or 11) long skip connections b/w encoder-decoder or should I still stick with a connection per layer - exactly as your provided snippet? <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"></p>
<p>Thanks a ton again!!! <img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=12" title=":hugs:" class="emoji" alt=":hugs:" loading="lazy" width="20" height="20"></p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="neel_g" data-post="7" data-topic="145408">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neel_g/48/42759_2.png" class="avatar"> neel_g:</div>
<blockquote>
<p>In that case, should I just add a few (say 7 or 11) long skip connections b/w encoder-decoder or should I still stick with a connection per layer - exactly as your provided snippet?</p>
</blockquote>
</aside>
<p>Hmm, that’s an interesting question and I would probably start with the “simpler” implementation.<br>
I.e. if your current code would be easier to implement by creating skip connections in each layer, maybe try it out first. Then you could try to add conditions where the skip connections should be used and compare it against the first approach.</p>