Hi all,
I was able to implement Resnet34 using the class as follows
class Resnet34(nn.Module):
def __init__(self, num_classes = 10):
super(Resnet34,self).__init__()
original_model = models.resnet34(pretrained=True)
self.features = nn.Sequential(*list(original_model.children())[:-1])
self.classifier = nn.Sequential(nn.Linear(512, num_classes))
def forward(self, x):
f = self.features(x)
f = f.view(f.size(0), -1)
y = self.classifier(f)
return f,y
But when i try similar approach for densenet.
class Densenet161(nn.Module):
def __init__(self, num_classes = 2):
super(Densenet161,self).__init__()
original_model = models.densenet161(pretrained=True)
self.features = nn.Sequential(*list(original_model.children())[:-1])
self.classifier = (nn.Linear(2208, num_classes))
def forward(self, x):
f = self.features(x)
f = f.view(f.size(0), -1)
y = self.classifier(f)
return f,y
I get the following error
RuntimeError: size mismatch at /py/conda-bld/pytorch_1493676237139/work/torch/lib/THC/generic/THCTensorMathBlas.cu:243
I am using the same code used in Pytorch transfer learning fine tuning example.
Please suggest a method to sort the issue