Pre-trained ResNet50 as featurizer and predictor

I wanted to access the intermediate outputs of ResNet50, so I created my own class:

class ResNetFeat(nn.Module):
  def __init__(self):
    super(ResNetFeat, self).__init__()
    resnet50_pretrained = models.resnet50(pretrained=True)
    conv_modules = list(resnet50_pretrained.children())[:-2] #all layers expect Adaptive AvgPool2D and FC
    pool_fc_modules = list(resnet50_pretrained.children())[-2:] #layers Adaptive AvgPool2D and FC
    self.model_conv = nn.Sequential(*conv_modules)
    self.pool_fc = nn.Sequential(*pool_fc_modules)
  
  def forward(self, image):
    feat = self.model_conv(image)
    print(feat.shape)
    pred = self.pool_fc(feat)
    return feat, pred

When I test this code:

x_image = torch.randn(1, 3, 224, 224)


featurizer = ResNetFeat()
feat, pred = featurizer(x_image)
print(feat.shape, pred.shape)

I get the following error –
RuntimeError: size mismatch, m1: [2048 x 1], m2: [2048 x 1000] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:940

your self.pool_fc(feat) contains two global operations: pooling, and an fc layer.

Sequential(
  (0): AvgPool2d(kernel_size=7, stride=1, padding=0)
  (1): Linear(in_features=2048, out_features=1000, bias=True)
)

The latter needs the pooled features to be flattened first.
Here something to work with:

import torch
import torchvision.models as models
import torch.nn as nn


class ResNetFeat(nn.Module):
    def __init__(self):
        super(ResNetFeat, self).__init__()
        resnet50_pretrained = models.resnet50(pretrained=True)
        conv_modules = list(resnet50_pretrained.children())[:-2]  # all layers expect Adaptive AvgPool2D and FC
        pool_modules = [list(resnet50_pretrained.children())[-2]]  # layers Adaptive AvgPool2D: there is only one.
        fc_modules = [list(resnet50_pretrained.children())[-1]]  # fc modules (1 only)
        self.model_conv = nn.Sequential(*conv_modules)
        self.pool = nn.Sequential(*pool_modules)
        self.fc = nn.Sequential(*fc_modules)

    def forward(self, image):
        feat2d = self.model_conv(image)
        print(feat2d.shape)
        featedpooled = self.pool(feat2d)
        # flatten the feature before feeding the fc.
        featdense = featedpooled.view(featedpooled.size(0), -1)
        pred = self.fc(featdense)
        return featdense, pred


if __name__ == "__main__":
    x_image = torch.randn(1, 3, 224, 224)
    featurizer = ResNetFeat()
    feat, pred = featurizer(x_image)
    print(feat.shape, pred.shape)

Output:

torch.Size([1, 2048, 7, 7])
torch.Size([1, 2048]) torch.Size([1, 1000])