Search code examples
testingdeep-learningconv-neural-networkprediction

Saving and Loading CNN model


I have a CNN model and I want to save it and load it for prediction in different tab. But I am confused whether the model.evulotion part is included in the part I will save. And I don't know if it would be better to use Model.checkpoint or model.save to save and load. Is there anyone have an idea ? Thank you in advance

I'm in dilemma about using both of them so I've use it.


Solution

  • Using model.eval() just tells the PyTorch model to use mean values for batch normalisation, and deactivates the dropout layers. You can save your model without using model.eval() as it will not affect the performance. While saving model, saving model's state dictionary is preferred. This can be done as shown:

    #declare class of model here
    model = NeuralNetwork()
    #add training code below
    ...
    #saving model, the model will be saved at the intermediateWeightPath location
    intermediateWeightPath = "./bestmodel.pth"
    torch.save(model.state_dict(), intermediateWeightPath)