Renet101 modifying conv layer to accept a concatenated tensor

Tanvir_Roshid: <p>Hi,</p>
<p>I’ve loaded a resnet101 model that is pretrained. I have locked the gradients for layers 1 to 3</p>
<p>I have successfully been training it for my task</p>
<p>I would like to try enhance its performance through a VQ-VAE and extract an embed layer and concatenate it inside layer 4 during the forward pass.</p>
<pre><code class="lang-auto">class ResNet101(nn.Module):
    def __init__(self, num_classes, weights=ResNet101_Weights.IMAGENET1K_V2,
                 dropout_rate=0.5):
        super(ResNet101, self).__init__()

        self.resnet = models.resnet101(weights=weights)
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # Freeze all other parameters
        for param in self.resnet.parameters():
            param.requires_grad = False

        # Unfreeze Layer 4
        for param in self.resnet.layer4.parameters():
            param.requires_grad = True

        self.resnet.layer4[0].conv1 = nn.Conv2d(1536, 512, kernel_size=(1, 1), stride=(1, 1), bias=False).to(self.device)  # Move the conv1 weight to GPU

        num_features = self.resnet.fc.in_features
        self.resnet.fc = nn.Sequential(
            nn.Linear(num_features, 1024),
            nn.BatchNorm1d(1024),
            nn.LeakyReLU(),
            nn.Dropout(dropout_rate),

            nn.Linear(1024, 512),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(),
            nn.Dropout(dropout_rate),
        )

        self.skip_connection = nn.Linear(num_features, 512)

        self.fc_final = nn.Linear(1024, num_classes)

    '''
        # x is torch.Size([16, 1024, 14, 14])
        # z is torch.Size([16, 512, 14, 14])
        # x is now torch.Size([16, 1536, 14, 14])
    '''
    def forward(self, x, z):
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x = self.resnet.relu(x)
        x = self.resnet.maxpool(x)

        x = self.resnet.layer1(x)
        x = self.resnet.layer2(x)
        x = self.resnet.layer3(x)

        # Concatenate z with the input
        x = torch.cat((x, z), dim=1).to(self.device)  # Move x to GPU

        x = self.resnet.layer4(x)

        x_avg = self.resnet.avgpool(x)
        x_avg = torch.flatten(x_avg, 1)

        x_fcl = self.resnet.fc(x_avg)
        x_skip = self.skip_connection(x_avg)
        x_fcl_skip = torch.cat((x_fcl, x_skip), dim=1)

        x_final = self.fc_final(x_fcl_skip)

        return x_final
</code></pre>
<p>At the moment, it is not working because:</p>
<pre><code class="lang-auto">RuntimeError                              Traceback (most recent call last)
Cell In[42], line 31
     28 classifier_optimizer.zero_grad()
     30 # Forward pass
---&gt; 31 vq_vae_output, classifier_output = vq_vae_classifier(images)
     33 # Compute the VQ-VAE loss, classifier loss and total loss
     34 vq_vae_loss, classifier_loss, total_loss = criterion(vq_vae_output, classifier_output, images, statuses)

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torch\nn\modules\module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

Cell In[29], line 15, in VQVAEClassifier.forward(self, images)
     11 z = self.vq_vae.encoder(images).to(device)
     13 images_normalized = normalize_batch(images, IMG_MEAN, IMG_STD)
---&gt; 15 predicted_statuses = self.classifier(images_normalized, z)
     17 return ( (images_reconstructed, commitment_loss, codebook_loss, perplexity), predicted_statuses )

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torch\nn\modules\module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\final\deep_learning\resnet101\resnet101.py:64, in ResNet101.forward(self, x, z)
     61 # Concatenate z with the input
     62 x = torch.cat((x, z), dim=1)
---&gt; 64 x = self.resnet.layer4(x)
     66 x_avg = self.resnet.avgpool(x)
     67 x_avg = torch.flatten(x_avg, 1)

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torch\nn\modules\module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torch\nn\modules\container.py:217, in Sequential.forward(self, input)
    215 def forward(self, input):
    216     for module in self:
--&gt; 217         input = module(input)
    218     return input

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torch\nn\modules\module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torchvision\models\resnet.py:158, in Bottleneck.forward(self, x)
    155 out = self.bn3(out)
    157 if self.downsample is not None:
--&gt; 158     identity = self.downsample(x)
    160 out += identity
    161 out = self.relu(out)

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torch\nn\modules\module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torch\nn\modules\container.py:217, in Sequential.forward(self, input)
    215 def forward(self, input):
    216     for module in self:
--&gt; 217         input = module(input)
    218     return input

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torch\nn\modules\module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torch\nn\modules\conv.py:463, in Conv2d.forward(self, input)
    462 def forward(self, input: Tensor) -&gt; Tensor:
--&gt; 463     return self._conv_forward(input, self.weight, self.bias)

File c:\Users\Tanvi\Desktop\comp5200m-msc-project\venv\Lib\site-packages\torch\nn\modules\conv.py:459, in Conv2d._conv_forward(self, input, weight, bias)
    455 if self.padding_mode != 'zeros':
    456     return F.conv2d(F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode),
    457                     weight, bias, self.stride,
    458                     _pair(0), self.dilation, self.groups)
--&gt; 459 return F.conv2d(input, weight, bias, self.stride,
    460                 self.padding, self.dilation, self.groups)

RuntimeError: Given groups=1, weight of size [2048, 1024, 1, 1], expected input[16, 1536, 14, 14] to have 1024 channels, but got 1536 channels instead
</code></pre>
<p>I’m not sure how to resolve this? Is it even possible for me to do</p>

ptrblck: <p>The error is raised in <code>self.resnet.layer4</code> since you are concatenating <code>x</code> with <code>z</code> increasing the channel dimension from the original <code>1024</code> to <code>1536</code>.<br>
You would thus need to replace the conv layer in <code>self.resnet.layer4</code> with a newly initialized one accepting <code>1536</code> channels and train it.</p>

Tanvir_Roshid: <aside class="quote no-group" data-username="Tanvir_Roshid" data-post="1" data-topic="183010">
<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/tanvir_roshid/48/61259_2.png" class="avatar"> Tanvir_Roshid:</div>
<blockquote>
<pre><code class="lang-auto">        self.resnet.layer4[0].conv1 = nn.Conv2d(1536, 512, kernel_size=(1, 1), stride=(1, 1), bias=False).to(self.device)  # Move the conv1 weight to GPU
</code></pre>
</blockquote>
</aside>
<p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, thank you for your reply.</p>
<p>I assumed I was with the line above?</p>

Tanvir_Roshid: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> I managed to get it working:</p>
<pre><code class="lang-auto">
import torch

import torch.nn as nn
import torchvision.models as models
from torchvision.models import ResNet101_Weights


class ResNet101(nn.Module):
    def __init__(self, num_classes, weights=ResNet101_Weights.IMAGENET1K_V2,
                 dropout_rate=0.5):
        super(ResNet101, self).__init__()
        self.resnet = models.resnet101(weights=weights)

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # Freeze all other parameters
        for param in self.resnet.parameters():
            param.requires_grad = False

        # Unfreeze Layer 4
        for param in self.resnet.layer4.parameters():
            param.requires_grad = True

        # Modify the convolutional layer in Layer 4
        self.resnet.layer4[0].conv1 = nn.Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)

        num_features = self.resnet.fc.in_features
        self.resnet.fc = nn.Sequential(
            nn.Linear(num_features, 1024),
            nn.BatchNorm1d(1024),
            nn.LeakyReLU(),
            nn.Dropout(dropout_rate),

            nn.Linear(1024, 512),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(),
            nn.Dropout(dropout_rate),
        )

        self.skip_connection = nn.Linear(num_features, 512)

        self.fc_final = nn.Linear(1024, num_classes)

    def forward(self, x, z=None):
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x = self.resnet.relu(x)
        x = self.resnet.maxpool(x)

        x = self.resnet.layer1(x)
        x = self.resnet.layer2(x)
        x = self.resnet.layer3(x)

        if z is not None:
            # Concatenate x and z along the channel dimension
            x_concat = torch.cat((x, z), dim=1).to(self.device)  # Move x to GPU
        else:
            x_concat = x

        x_concat = self.resnet.layer4[0].conv1(x_concat)
        x_concat = self.resnet.layer4[0].bn1(x_concat)

        x_concat = self.resnet.layer4[0].conv2(x_concat)
        x_concat = self.resnet.layer4[0].bn2(x_concat)

        x_concat = self.resnet.layer4[0].conv3(x_concat)
        x_concat = self.resnet.layer4[0].bn3(x_concat)

        x_concat = self.resnet.layer4[0].relu(x_concat)

        x_concat = self.resnet.layer4[0].downsample(x)  # Apply downsample to the output of layer 3

        x_concat = self.resnet.layer4[1](x_concat)
        x_concat = self.resnet.layer4[2](x_concat)

        x_avg = self.resnet.avgpool(x_concat)
        x_avg = torch.flatten(x_avg, 1)

        x_fcl = self.resnet.fc(x_avg)
        x_skip = self.skip_connection(x_avg)
        x_fcl_skip = torch.cat((x_fcl, x_skip), dim=1)

        x_final = self.fc_final(x_fcl_skip)

        return x_final
</code></pre>
<p>However, I am a novice, I need to be sure that this code is correct, i.e., it still maintains the original <code>resnet101</code> forward pass with only the <strong>modification</strong> that I intended.</p>
<p>I edited it slightly, i.e., both x and z are <code>[16, 1024, 14, 14]</code></p>

Tanvir_Roshid: <p>Based on <a href="https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py#L143" rel="noopener nofollow ugc">here</a></p>
<p>I would argue the correct forward statement is:</p>
<pre><code class="lang-auto">    def forward(self, x, z=None):
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x = self.resnet.relu(x)
        x = self.resnet.maxpool(x)

        x = self.resnet.layer1(x)
        x = self.resnet.layer2(x)
        x = self.resnet.layer3(x)

        if z is not None:
            # Concatenate x and z along the channel dimension
            x_concat = torch.cat((x, z), dim=1).to(self.device)  # Move x to GPU
        else:
            x_concat = x

        x_concat = self.resnet.layer4[0].conv1(x_concat)
        x_concat = self.resnet.layer4[0].bn1(x_concat)
        x_concat = self.resnet.layer4[0].relu(x_concat)

        x_concat = self.resnet.layer4[0].conv2(x_concat)
        x_concat = self.resnet.layer4[0].bn2(x_concat)
        x_concat = self.resnet.layer4[0].relu(x_concat)

        x_concat = self.resnet.layer4[0].conv3(x_concat)
        x_concat = self.resnet.layer4[0].bn3(x_concat)

        x = self.resnet.layer4[0].downsample(x)  # Apply downsample to the output of layer 3
        x_concat += x
        x_concat = self.resnet.layer4[0].relu(x_concat)

        x_concat = self.resnet.layer4[1](x_concat)
        x_concat = self.resnet.layer4[2](x_concat)

        x_avg = self.resnet.avgpool(x_concat)
        x_avg = torch.flatten(x_avg, 1)

        x_fcl = self.resnet.fc(x_avg)
        x_skip = self.skip_connection(x_avg)
        x_fcl_skip = torch.cat((x_fcl, x_skip), dim=1)

        x_final = self.fc_final(x_fcl_skip)

        return x_final
</code></pre>