I have a pretrained model that I use for semantic segmentation. It predicts one of 50 classes for each pixel.
I want to further train it on my own dataset, but for predicting only 1 of 10 possible classes (a subset of the 50 classes). The other 40 classes don’t appear in the new dataset.
Obviously, I can just ignore the other classes, but this way I lose much of the accuracy already obtained in the pretrained model.
I want to modify the last layer of the model, and change the saved pretrained model to fit it.
The last layers of the model are:
num_classes = 50
self.last_layer = nn.sequential(
    nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1, bias=False),
    nn.ReLU(inplace=True),
    nn.Conv2d(512, num_classes, kernel_size=1))

The forward pass does that:
x = .......
x = self.last_layer(x)
x == nn.functional.log_softmax(x, dim=1)

And evaluates loss like that:
loss = my_criterion(x, gt)  # (my_criterion = nn.NLLLoss(ignore_index=-1) )

The weights are loaded like that:
model.load_state_dict(torch.load(weights_path, map_location=lambda s, l: s), strict=False)

I really don’t want to retrain the whole modified model.
I understand that another possibility is to just keep the model and take the predictions, but before running the softmax - to zero all of the predictions of the other classes. But it feels like keeping parts of the model that I don’t need and making the training more difficult.This text will be hidden