How can I load an exported model (or .pth file) for use on a different dataset (with a different number of classes)?

[zache (Zach Eberhart)]

I trained a model and am now trying to use it on a different, but similar, dataset. However, I am getting an error because there is a different amount of classes on the new dataset. How can I load it? (as I would any other pre-trained model)

Thanks!

[marcmuc (Marc P. Rostock)]

You don’t specify which type of model you have, maybe this thread below will help. Otherwise look for other threads or the docs for “custom head”, you have to basically remove the classifier head again and replace it with a new head with the desired num of classes.

New custom head in trained model

You have to dig a bit into the code base for this. The model created by create_cnn is a sequential model with two things: the body and the head. If you want to keep the pretrained body, you’d find it in model[0], the head is model[1]. If you want a new CNN with less classes, you can create your new model with learn1 = create_cnn(data1,...) then load the weights from the previous head into the new one with learn1.model[0].load_state_dict(learn.model[0].state_dict())

[zache (Zach Eberhart)]

Yup, that did the trick. Thanks!