6
$\begingroup$

So I am currently training some deep learning models for some basic classification problems, and I am trying to figure out if it is possible to change the output size of the model in case I want to retrain my model with a different datasets in which the number of classes is different.

I have seen a lot of "retraining" posts but I don't get totally how is this internally occurring too. Are all the weights being updated with the new dataset except for the ones in the last layer that connects the previous last connection to a new output layer with a different amount of neurons. I also saw that is sometimes common to retrain only the last layer. How is this done? And how would this be coded for a model in Pytro

$\endgroup$

1 Answer 1

6
$\begingroup$

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()
$\endgroup$
2
  • $\begingroup$ Amazing, thanks a lot for such an informative answer! I didn't new that the layers of the model could be accessed as you did in model[-1] = nn.Linear(model[-1].in_features, n_new_classes). $\endgroup$ May 17 at 21:48
  • 1
    $\begingroup$ My pleasure, glad it was helpful. Yeah the indexing is a convenient way of accessing models. When debugging I also use model[:2](input), model[:3](input), etc to progressively check the output at different levels of the model for a given input. $\endgroup$ May 17 at 21:59

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.