Retraining a UNet model

[rirase]

Thanks to everyone for this library, it’s great.

I have a question: I trained a UNet model with 5 labels for 5 bones, but now I need a segmentation for 7 bones, so I need to include two more. I am going to use the same dataset. Can I retrain my model, which is already working well, or do I need to start from scratch? When I try the demo to train, it gives me an error saying that the keys of the new UNet don’t match.
Thanks for your help.

Best regards,
Ricar

[KumoLiu]

Hi @rirase, you can either retrain the model or start from scratch. If you want retrain, you need try to adjust the output layer as your output class has changed such as (you need to check what's your final layer):

# Change the final layer to 7 classes:
model.final_conv = nn.Conv2d(in_channels, 7, kernel_size=1)

Hope it helps, thanks.

[ericspod]

This is a good way to get started with repurposing a network. If the weight loading doesn't work it's probably only the final_conv weights that have different shape, so omit these from the state dictionary you're loading or set the strict argument of load_state_dict to False. With your weights loaded, you can train the new model by either optimising the whole network during training or only just the new layer, you'll probably have to figure out experimentally what works better.