Pytorch 保存中途训练的模型

一般看到的是这个：

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')