一般看到的是这个:
torch.save(net.state_dict(), 'model.pth')
但这个只保存了模型本身的参数,能保证可以load保存的模型进行测试,但是并不方便恢复训练。
实际上还需要保存优化器的状态,以及一些其他有助于恢复训练的信息(如loss等):
def save_checkpoint(epoch, epochs_since_improvement, model, optimizer, loss, is_best):
state = {'epoch': epoch,
'epochs_since_improvement': epochs_since_improvement,
'loss': loss,
'model': model,
'optimizer': optimizer}
filename = 'checkpoint_' + str(epoch) + '_' + str(loss) + '.tar'
torch.save(state, filename)
if is_best:
torch.save(state, 'BEST_checkpoint.tar')