保存和加载整个模型
torch.save(net, 'model.pkl')
net = torch.load('model.pkl')
仅保存和加载模型参数
state = {
'net': net.state_dict(),
'acc': acc,
'epoch': epoch,
}
torch.save(state, 'params.pkl')
torch.save(net.state_dict(), 'params.pkl')
net.load_state_dict(torch.load('params.pkl'))