I am trying to create a custom nn module that wraps a pre-trained densenet121, so that I can access activations at the end of the convolutional part of the network.
As a starting point, I have just done the following:
import torch from torchvision import models import torch.nn as nn class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() self.model = models.densenet121(pretrained=True) self.features = self.model.features self.classifier = self.model.classifier def forward(self, x): x = self.features(x) x = self.classifier(x) return x example = torch.ones((1,3,224,224)) model = MyModel() output = model.forward(example)
However, I get a mismatch error for line
x = self.classifier(x) of the forward method:
RuntimeError: size mismatch, m1: [7168 x 7], m2: [1024 x 1000] at ../aten/src/TH/generic/THTensorMath.cpp:961
Also, as far as I can tell, densenet121 expects (batch_size, 1024) inputs to its classifier part, however when I print the size of x after
x = self.features(x) in the forward method, it is (1, 1024, 7, 7)?
Thank you for the help, I am still very new to pytorch.