Things on this page are fragmentary and immature notes/thoughts of the author. Please read with your own judgement!
It is common to customize a pretrained model by delete the output layer or replace it to the output layer that suits your use case. There are several ways to achieve this in PyTorch.
Replace the Fully Connected Layer with an Identity Layer¶
Define an identity layer.
Replace the fully connected layer with an Identity Layer (using ResNet18 as an example).
model = models.resnet18(pretrained=False)
model.fc = Identity()import torch
from torchvision import models
class Identity(torch.nn.Module):
def __init__(self):
super(Identity, self).__init__()
def forward(self, x):
return xExtract Layers and Chain Them¶
model.classifier = nn.Sequential(*[model.classifier[i] for i in range(4)])
print(model.classifier)model.classifier = nn.Sequential(*list(model.classifier.children())[:-3])Inference using the customized model.
x = torch.randn(1, 3, 224, 224)
output = model(x)
print(output.shape)torch.Size([1, 512])
Using Forward Hooks¶
Please refer to How can l load my best model as a feature extractor/evaluator? for more details.