model_ft = models.resnet18(pretrained=True)
num_ftrs = model_ft.fc.in_features
# Here the size of each output sample is set to 2.
# Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).
model_ft.fc = nn.Linear(num_ftrs, 2)
However, I cannot know how to insert layers before the transferred model.
And, I am not sure method of adding layers to the transferred model after, the above code.
Hi,
you could create a new class for your model, which contains your layers before the model + the transfer model with a changed last layer (how you did int line 2 of your example) + all your additional layers at the end. In your forward method you now call your ‘before layers’, then your transferred model and last your ‘after layers’.
class MyModel(nn.Model):
def __init__(self, out_features, use_pretrained):
self.pre_layers = ...
self.backbone = models.vgg16(pretrained=use_pretrained)
self.backbone[6] = nn.Linear(in_features=4096, out_features=out_features)
self.after_layers = ...
def forward(self, x):
x = self.pre_layers(x)
x = self.backbone(x)
x = self.after_layers(x)
return x