So if I understand you correctly in my case I wouldn’t need this line: model = TheModelClass(*args, **kwargs) ?
So just retrieving the model like this would be enough?
model = torch.load(PATH)
model.load_state_dict(torch.load(PATH))
model.eval()




 ptrblck:

In that case i think just swapping the last layer in the classifier might be sufficient:
model = models.segmentation.fcn_resnet101(pretrained=True)

nb_classes = 2
model.classifier[4] = nn.Conv2d(
    in_channels=512,
    out_channels=nb_classes,
    kernel_size=1,
    stride=1)



Also, I am a little confused with what you mentioned here, because I already update/ or create a new layer in the finetuning process when I initialize the model like in the function below. So wouldn’t doing this again when retrieving the saved finetuned model be redundant?
def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):

    # Initialize these variables which will be set in this if statement. Each of these

    #   variables is model specific.

    model_ft = None

    input_size = 0

    if model_name == "resnet":

        """ FCN_resnet101

        """

        model_ft = models.segmentation.fcn_resnet101(pretrained=use_pretrained)

        set_parameter_requires_grad(model_ft, feature_extract)

        in_chnls = model_ft.classifier[4].in_channels

        model_ft.classifier[4] = nn.Conv2d(in_chnls, num_classes, 1, 1)

        input_size = 750, 1000 

    else:

        print("Invalid model name, exiting...")

        exit()

    return model_ft, input_size

I then use model_ft as in input in my train_model function which returns the fine tuned model
# Initialize the model for this run

model_ft, input_size = initialize_model(model_name, num_classes, feature_extract, use_pretrained=True)

# Train and evaluate

finetuned_model_ft, hist = train_model(model_ft, dataloaders_dict, criterion, optimizer_ft, num_epochs=num_epochs, is_inception=(model_name=="inception"))

So the next step would be then to save the new finetuned_model_ft like this:
torch.save(finetuned_model_ft.state_dict(), path/to/myModel)
and then retrieve it to be used for inference like this:
my_fine_tuned_model = torch.load(path/to/myModel)
my_fine_tuned_model.load_state_dict(torch.load(path/to/myModel))
my_fine_tuned_model.eval()

Is this correct?