Saving a model in pytorch

1

Hi I want to save my model during the training. My model has 9 conv layers with batch norm and softmax as activation function. I use this method to save my training modeel but when I resumed my training I sensed that it resumed from begining.

state = {
    'epoch': epoch,
    'state_dict': model.state_dict(),
    'optimizer': optimizer.state_dict(),
}
    torch.save(state, '/home/superblock/state_train.pt')

  state = torch.load('/home/superblock/state_train.pt')
  model.load_state_dict(state['state_dict'])
  optimizer.load_state_dict(state['optimizer'])

what should I do?
does it need model.eval()? I read some where that for resuming training we dont use that
and I have another question… It’s better to save when loss is fewer in validation data or train data?

2

model.eval() changes the behaviour of a few layers like batchnorm and dropout. So if you want to infer using your model then put the model in .eval() mode or if you want to continue training then put the model into training using model.train().
Save at the least validation loss

1 Like