path_model = 'model.pkl'
path_state_dict = 'model_state_dict.pkl'
1、保存(序列化)
# obj:对象
# f:输出路径
torch.save(obj,f)
2、保存模型参数 # 数据以字典形式存储
# 获取模型的可学习参数
state_dict = net.state_dict()
# 将参数保存
torch.save(state_dict,f)
3、加载(反序列化)
# f:文件路径
@当使用gpu时用map_location
# map_location:指定存放位置(cpu或gpu)
net = torch.load(f,map_location)
4、加载模型参数
# 加载数据
state_dict_load = torch.load(f,map_location)
# 将数据加载至模型
net.load_state_dict(state_dict_load)
5、断点续训练/定时保存模型及优化器等训练好的数据
@网络模型和优化器随学习而不断变化
# 需要保存的数据
checkpoint = {
# 模型数据
'model_state_dict':net.state_dict(),
# 优化器数据
'optimizer_state_dict':optimizer.state_dict(),
# epoch
'epoch':epoch
}
path = 'checkpoint_{}_epoch.pkl'.format(epoch)
torch.save(checkpoint,path)
# 加载
checkpoint = torch.load(path)
net.load_state_dict(checkpoint[model_state_dict'])
opimzer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch']