torch中国模型保存与加载

pytorch保存模型等相关参数，利用torch.save()，以及读取保存之后的文件

假设网络为model = Net(), optimizer = optim.Adam(model.parameters(), lr=args.lr), 假设在某个epoch，我们要保存模型参数，优化器参数以及epoch

一、保存模型

1. 先建立一个字典，保存三个参数：

state = {‘net':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}

2. 调用torch.save():

torch.save(state, dir)

其中dir表示保存文件的绝对路径+保存文件名，如’/home/qinying/Desktop/modelpara.pth’

二、加载恢复

checkpoint = torch.load(dir)

model.load_state_dict(checkpoint['net'])

optimizer.load_state_dict(checkpoint['optimizer'])

start_epoch = checkpoint['epoch'] + 1