loading saved model doesn't behave as expected when finetuning it

489 views Asked by At

I trained a pytorch model the accuracy at end of the first epoch is 20% and the loss value is 3.8 . I trained it until the loss is 3.2 and accuracy is around 50% and save it like this:

 torch.save(model.state_dict(), 'model.pth')

Then I load it like this:

model = Perceptron(input_size, output_size)
model.load_state_dict(torch.load("model.pth"))
model.to(device, dtype=torch.double)

When I'm starting to fine-tune it using the same task, same optimizer, and learning rate, I expect the loss starts at 3.2 and accuracy to be 50% but it looks like the model is rolling back and starts from a loss value of 3.8 and accuracy of 20% again. Is something wrong with my code or there's something that I don't understand about the fine-tuning model?

1

There are 1 answers

0
ki-ljl On

First, because you need to fine-tune, the optimizer also needs to be saved:

state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict()}
torch.save(state, "model.pth")

And then:

model.load_state_dict(torch.load("model.pth")['model])
optimizer.load_state_dict(torch.load("model.pth")['optimizer'])

Second, you need to set the random seed at the very beginning of the code

def setup_seed(seed):
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

setup_seed(20)