How to split a model, add a new layer and combine all the models

praveenck06: <p>In one of my use cases, I need to split trained models and add a custom layer in between to perform some calculations.</p>
<p>I have tried as follows</p>
<pre><code class="lang-auto">vgg_model = models.vgg11(pretrained=True)

class CustomLayer(nn.Module):
    def __init__(self): 
        super().__init__()

        
    def forward(self, input_features):
        input_features = input_features*0.5 # some calculations
        return input_features  


class NetworkWithCustomeLayer(nn.Module):
    def __init__(self, trained_model, layer_number):
        super(NetworkWithCustomeLayer, self).__init__()
        self.initial_layers =  nn.Sequential(*list(trained_model.children())[:layer_number])
        self.custom_layer = CustomLayer()
        self.last_layers =  nn.Sequential(*list(trained_model.children())[layer_number:])
    def forward(self, input):
        print("input shape",input.shape)
        x_1 = self.initial_layers(input)
        print("x1 shape",x_1.shape)
        x_c = self.custom_layer(x_1)
        print("xc shape",x_c.shape)
        x_2 = self.last_layers(x_c)
        print("x2 shape",x_2.shape)
        return x_2
modified_vgg = NetworkWithCustomeLayer(vgg_model, 1)
modified_vgg.to(device)
</code></pre>
<p>this is how modified_vgg looks</p>
<pre><code class="lang-auto">NetworkWithCustomeLayer(
  (initial_layers): Sequential(
    (0): Sequential(
      (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (1): ReLU(inplace=True)
      (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (4): ReLU(inplace=True)
      (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (6): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (7): ReLU(inplace=True)
      (8): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (9): ReLU(inplace=True)
      (10): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (11): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (12): ReLU(inplace=True)
      (13): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (14): ReLU(inplace=True)
      (15): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (16): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (17): ReLU(inplace=True)
      (18): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (19): ReLU(inplace=True)
      (20): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
  )
  (custom_layer): CustomLayer()
  (last_layers): Sequential(
    (0): AdaptiveAvgPool2d(output_size=(7, 7))
    (1): Sequential(
      (0): Linear(in_features=25088, out_features=4096, bias=True)
      (1): ReLU(inplace=True)
      (2): Dropout(p=0.5, inplace=False)
      (3): Linear(in_features=4096, out_features=4096, bias=True)
      (4): ReLU(inplace=True)
      (5): Dropout(p=0.5, inplace=False)
      (6): Linear(in_features=4096, out_features=1000, bias=True)
    )
  )
)
</code></pre>
<p>but when evaluting the model, I am getting following error:<br>
<strong>RuntimeError: mat1 and mat2 shapes cannot be multiplied (114688x7 and 25088x4096)</strong></p>
<p>below is the stack trace of the error:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_24462/3433225722.py in &lt;module&gt;
     13     print('Accuracy of the network on test images: %d %%' % (
     14         100 * correct / total))
---&gt; 15 evaluate_model(modified_vgg)

/tmp/ipykernel_24462/3433225722.py in evaluate_model(model)
      6             images = images.to(device)
      7             labels = labels.to(device)
----&gt; 8             output=  model(images)
      9             _, predicted = torch.max(output, 1)
     10             total += labels.size(0)

/anaconda/envs/py38_pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/tmp/ipykernel_24462/645674472.py in forward(self, input)
     23         x_c = self.custom_layer(x_1)
     24         print("xc shape",x_c.shape)
---&gt; 25         x_2 = self.last_layers(x_c)
     26         print("x2 shape",x_2.shape)
     27         return x_2

/anaconda/envs/py38_pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/anaconda/envs/py38_pytorch/lib/python3.8/site-packages/torch/nn/modules/container.py in forward(self, input)
    137     def forward(self, input):
    138         for module in self:
--&gt; 139             input = module(input)
    140         return input
    141 

/anaconda/envs/py38_pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/anaconda/envs/py38_pytorch/lib/python3.8/site-packages/torch/nn/modules/container.py in forward(self, input)
    137     def forward(self, input):
    138         for module in self:
--&gt; 139             input = module(input)
    140         return input
    141 

/anaconda/envs/py38_pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/anaconda/envs/py38_pytorch/lib/python3.8/site-packages/torch/nn/modules/linear.py in forward(self, input)
     94 
     95     def forward(self, input: Tensor) -&gt; Tensor:
---&gt; 96         return F.linear(input, self.weight, self.bias)
     97 
     98     def extra_repr(self) -&gt; str:

/anaconda/envs/py38_pytorch/lib/python3.8/site-packages/torch/nn/functional.py in linear(input, weight, bias)
   1845     if has_torch_function_variadic(input, weight):
   1846         return handle_torch_function(linear, (input, weight), input, weight, bias=bias)
-&gt; 1847     return torch._C._nn.linear(input, weight, bias)
   1848 
   1849 

RuntimeError: mat1 and mat2 shapes cannot be multiplied (114688x7 and 25088x4096)

</code></pre>
<p>when we see the output of the custom layer its same as the inital_layers, but getting shape mismatch error for final layers. dont know why the data shape is getting changed</p>

ptrblck: <p>Rewrapping the modules in an <code>nn.Sequential</code> block can easily break, since you would miss all functional API calls from the original <code>forward</code> method and will thus only work if the layers are initialized and executed sequentially.<br>
For <code>VGG11</code> you would be missing the <code>torch.flatten</code> operation from <a href="https://github.com/pytorch/vision/blob/f3365262d1ba7233ea70b4020d559f1dcc3f0612/torchvision/models/vgg.py#L55-L60">here</a>, which would create the shape mismatch.<br>
The safest way would be to create a custom model, add the custom layer, and override the <code>forward</code>.<br>
Alternatively, you could also replace an already used layer (e.g. the last <code>nn.MaxPool2d</code> layer) with an <code>nn.Sequential</code> container containing itself and the new custom layer.</p>

praveenck06: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>.<br>
separating fc layer and adding flatten before it helped.</p>
<pre><code class="lang-auto">class CustomLayer(nn.Module):
    def __init__(self): 
        super().__init__()

        
    def forward(self, input_features):
        input_features = input_features*0.5 # some calculations
        return input_features  


class NetworkWithCustomeLayer(nn.Module):
    def __init__(self, trained_model, layer_number):
        super(NetworkWithCustomeLayer, self).__init__()
        self.initial_layers =  nn.Sequential(*list(trained_model.children())[:layer_number])
        self.custom_layer = CustomLayer()
        self.last_layers =  nn.Sequential(*list(trained_model.children())[layer_number:-1])
        self.fc = list(vgg_model.children())[-1]
    def forward(self, input):
        print("input shape",input.shape)
        x_1 = self.initial_layers(input)
        print("x1 shape",x_1.shape)
        x_c = self.custom_layer(x_1)
        print("xc shape",x_c.shape)
        x_2 = self.last_layers(x_c)
        print("x2 shape",x_2.shape)
        x_3 = torch.flatten(x_2, 1)
        print("x3 shape",x_3.shape)
        x_4 = self.fc(x_3) 
        return x_4
modified_vgg = NetworkWithCustomeLayer(vgg_model, 1)
modified_vgg.to(device)
</code></pre>