In general, when you want to use an existing net for a new task that has different output requirements, you would usually remove the output layer and replace it with a new one.
Freeze all of the layers except the new one, and train the model for a bit in order to converge on sensible weights for the newly-added layer.
If you find that validation performance isn't as required, you could unfreeze the layer before, and gently nudge it towards better weights by using a low learning rate.
The less data you have and the more similar the tasks are, the less you'll need to tweak the existing layers, so you might find it adequate to just tune the newly-added layer, keeping the rest frozen.
You could drop some of the top layers if performance still isn't good and your dataset is limited (for similar tasks, the top layers are less useful than the bottom layers). If you have enough data, you could instead replace the top layers or make the net deeper.
The code for swapping out the last layer and freezing the others is along the lines of:
...
# Freeze the original layers
for param in model.parameters():
param.requires_grad = False
# Replace the last layer with a new one that outputs "n_new_classes"
model[-1] = nn.Linear(model[-1].in_features, n_new_classes)
#Learn weights for the new layer
optimizer = optim.SGD(model.parameters(), lr=1e-3)
for epoch in range(3):
for inputs, labels in train_loader:
optimizer.zero_grad()
outputs = model(inputs)
loss = nn.CrossEntropyLoss()(outputs, labels)
loss.backward()
optimizer.step()
If necessary, also tweak the preceding layer(s).
#Unfreeze the second-to-last layer
for param in model[-2].parameters():
param.requires_grad = True
#Tweak the trainable weights using a smaller learning rate
optimizer = optim.SGD(model.parameters(), lr=1e-4)
for epoch in range(1):
for inputs, labels in train_loader:
optimizer.zero_grad()
outputs = model(inputs)
loss = nn.CrossEntropyLoss()(outputs, labels)
loss.backward()
optimizer.step()
