Want to output intermediate layers from pretrained Resnet 18

I want to return the output from the layer 4, avgpool layer and the layer before the fc layer from resnet 18. For that I have written the following code which is showing error as avgpool and fc are being called twice (Once when I define self.avgpool and self.fc, again when I call these in forward method). Thing is, when I create self.avgpool or self.fc, these automatically adds a avgpool and fc layer at the end of resnet even if I don’t call those in the forward method. How can I do this?

params = {
    "model": "resnet18",
    "device": "cuda",
    "lr": 0.001,
    "batch_size": 64,
    "num_workers": 4,
    "epochs": 10,
    "num_classes": 1
}
class pretrainedModelBlock(nn.Module):
  def __init__(self, model_name = params['model'], classes = params["num_classes"], addDecoder = False):
    super(pretrainedModelBlock, self).__init__()
    self.model_name = model_name
    self.classes = classes
    self.addDecoder = addDecoder
    self.custom_model = self.create_model()
    self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
    self.fc = nn.Linear(in_features=512, out_features= self.classes, bias=True)

  def create_model(self):
    pretrained_model = getattr(models, self.model_name)(pretrained = True)
    pretrained_model.avgpool = nn.Identity()
    pretrained_model.fc = nn.Identity()
    return pretrained_model

  def forward(self, x):
    resnet_features = self.custom_model(x)
    pool = self.avgpool(resnet_features)
    flat = torch.flatten(pool, 1)
    out = self.fc(resnet_features)


    if self.addDecoder:
      return resnet_features, pool, flat, out
    
    else:
      return out

I think the issue has been solved in this post.

Alternatively, if you want to remove the last layers and keep the layers till you want link.

Hi @mxahan, thanks for the link. Btw, I replaced the avgpool and fc layers with nn.Identity(). Before the avgpool layer, the dimension of the tensor should be torch.Size([64, 512, 1, 1]) (My input size was (64,3,32,32)). But, even though I replaced both avgpool and fc with nn.Identity(), the result I’m getting has shape torch.Size([64,512]). I think the avgpool layer is still being activated somehow.

I tried manually adding the avgpool and fc layers with self.avgpool and self.fc after replacing the avgpool and fc from resnet with nn.Identity(). But this is not working (The commented out lines)

class pretrainedModelBlock(nn.Module):
    def __init__(
        self,
        model_name=params["model"],
        classes=params["num_classes"],
        addDecoder=False,
    ):
        super(pretrainedModelBlock, self).__init__()
        self.model_name = model_name
        self.classes = classes
        self.addDecoder = addDecoder
        self.pretrained_model = getattr(models, self.model_name)(pretrained=True)
        self.pretrained_model.avgpool = nn.Identity()
        self.pretrained_model.fc = nn.Identity()

        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
        self.fc = nn.Linear(in_features=512, out_features=self.classes, bias=True)

    def forward(self, x):
        resnet_features = self.pretrained_model(x)
        # pool = self.avgpool(resnet_features)
        # out = self.fc(pool)

        return resnet_features

Here’s my code, I got the output shape of [Bs, 512, 1, 1]

import torch
import torchvision

rnet = torchvision.models.resnet18(pretrained=True)

activation = {}
def get_activation(name):
    def hook(model, input, output):
        activation[name] = output.detach()
    return hook

rnet.layer4[0].conv2.register_forward_hook(get_activation('con2_layer'))

Bs = 16
inputs=  torch.rand([Bs, 3, 32, 32])
oupt =  rnet(inputs)

print(activation['con2_layer'].shape)

Sample Output

>>>torch.Size([16, 512, 1, 1])

As an alternative, you can remove the last two layers and create a new model as below

model = nn.Sequential(*list(rnet.children())[:-2])
oupt = model(inputs)
print(oupt.shape)

Output

torch.Size([16, 512, 1, 1])
1 Like