class ResNet(nn.Module):
def __init__(self, output_features, fine_tuning=False):
super(ResNet, self).__init__()
self.resnet152 = tv.models.resnet152(pretrained=True)
#freezing the feature extraction layers
for param in self.resnet152.parameters():
param.requires_grad = fine_tuning
#self.features = self.resnet152.features
self.num_fts = 512
self.output_features = output_features
# Linear layer goes from 512 to 1024
self.classifier = nn.Linear(self.num_fts, self.output_features)
nn.init.xavier_uniform_(self.classifier.weight)
self.tanh = nn.Tanh()
def forward(self, x):
h = self.resnet152(x)
print('h: ',h.shape)
return h
image_model_resnet152=ResNet(output_features=10).to(device)
image_model_resnet152
Here, after printing the image_model_resnet152
, I get:
Here, what is the difference between (avgpool): Linear(in_features=2048)
and (classifier): Linear(in_features=512)
?
I am implementing an image captioning model, so which in_features
should I take for an image?