How to get output from intermediate layer in nn.Sequential and reshape it before passing to next layer?

username: Sourabh_Patil
cooked: <p>Hello!</p>
<p>I have a pretrained model whose weights I wish to use. It has the last sequential layer as follows:</p>
<pre><code>    self.output_layer = Sequential(BatchNorm2d(512), 
                                   Dropout(drop_ratio),
                                   Flatten(),
                                   Linear(512 * 7 * 7, 512),
                                   BatchNorm1d(512))
</code></pre>
<p>The linear layer gives me output of shape 1 x 512. But in place of BatchNorm1d, I need to use Batchnorm2d as it is supported for conversion on the edge device on which I am working, and Batchnorm1d is not supported. So, I can change Batchnorm1d to Batchnorm2d and load the weights of pretrained model. But, Batchnorm2d expects a 4D tensor and Lienar layer will give me 2D tensor. So how to solve this? I suppose we can not add a reshape module in between Linear and Batchnorm2d as while loading the weights it will throw as error. Can we get output after Linear and then reshape it using a custom reshape module and then put it to BatchNorm2d again in the same sequential block? How to do it if its doable? Or is there any other workaround for this?</p>
<p>Thanks!</p>

username: ptrblck
cooked: <p>You could try to <code>insert</code> a custom layer to reshape the activation after loading the <code>state_dict</code>:</p>
<pre><code class="lang-python"># create reference
output_layer = nn.Sequential(nn.BatchNorm2d(512), 
                             nn.Dropout(0.5),
                             nn.Flatten(),
                             nn.Linear(512 * 7 * 7, 512),
                             nn.BatchNorm1d(512))

sd = copy.deepcopy(output_layer.state_dict())

# restore
layer = nn.Sequential(nn.BatchNorm2d(512), 
                      nn.Dropout(0.5),
                      nn.Flatten(),
                      nn.Linear(512 * 7 * 7, 512),
                      nn.BatchNorm2d(512))

layer.load_state_dict(sd)
# &lt;All keys matched successfully&gt;

x = torch.randn(2, 512, 7, 7)
out = layer(x)
# ValueError: expected 4D input (got 2D input)

class Reshape(nn.Module):
    def __init__(self, shape):
        super().__init__()
        self.shape = shape
        
    def forward(self, x):
        x = x.reshape(self.shape)
        return x

layer.insert(-1, Reshape((-1, 512, 1, 1)))
out = layer(x)
print(out.shape)
# torch.Size([2, 512, 1, 1])
</code></pre>

username: Sourabh_Patil
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> ,</p>
<p>Thank you for the apt solution. I did some other work around which is little lengthy and naive I feel.</p>
<p>Created a new ReshapeModule</p>
<p>class ReshapeModule(Module):<br>
def forward(self, x):<br>
x = x.view(1, 512, 1, 1)<br>
return x</p>
<p>In model’s forward method,<br>
…<br>
x = self.output_layer<a>:-1</a><br>
x = self.to_2d_bn(x)<br>
x = self.output_layer<a>-1:</a><br>
…</p>
<p>where,</p>
<p>self.output_layer = Sequential(BatchNorm2d(512),<br>
Dropout(drop_ratio),<br>
Flatten(),<br>
Linear(512 * 7 * 7, 512),<br>
BatchNorm2d(512))</p>
<p>Thanks,<br>
Sourabh</p>

username: ptrblck
cooked: <p>Your approach is also totally fine.</p>