欢迎来到尧图网

客户服务 关于我们

您的位置:首页 > 教育 > 锐评 > pytorch持久化

pytorch持久化

2025/2/24 0:48:32 来源:https://blog.csdn.net/liugan528/article/details/142099279  浏览:    关键词:pytorch持久化

在pytorch中以下对象可以持久化到硬盘,并能通过相应的方法加载到内存:Tensor、Variable、nn.Module、Optimizer

本质上上述信息最终都是保存成Tensor,Tensor的保存和加载是使用torch.savetorch.load完成的。在save、load时可以指定使用的pickle模块,在load时还可将GPU tensor映射到CPU或其它GPU上。

torch.save、torch.load

import torch# 1.变量的保存与加载
a = torch.Tensor(3,4)
print(a.get_device())  # -1,代表CPU
if torch.cuda.is_available():a = a.cuda()   # 把a转换为GPU0上的tensortorch.save(a, 'a.pth')# 加载为b,存储于GPU0上,因为保存时tensor就在GPU0上b = torch.load('a.pth')print(b.get_device())   # 0# 加载为d,存储于GPU0上c = torch.load('a.pth', map_location={'cuda:0':'cuda:0'})print(c.get_device())   # 0# 加载为c,存储于CPU上d = torch.load('a.pth', map_location=lambda storage, loc: storage)print(d.get_device())   # -1,代表CPU# 加载为d,存储于CPU上e = torch.load('a.pth', map_location={'cuda:0':'cpu'})print(e.get_device())   # -1,代表CPU# 2.模型的保存与加载
from torchvision.models import resnet18
model = resnet18()
torch.save(model.state_dict(), 'resnet18.pth')
model.load_state_dict(torch.load('resnet18.pth'))# 3.优化器的保存与加载
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
torch.save(optimizer.state_dict(), 'optimizer.pth')
optimizer.load_state_dict(torch.load('optimizer.pth'))# 4.保存模型和优化器的所有参数
all_data = dict(optimizer=optimizer.state_dict(), model=model.state_dict(), info=u'模型和优化器的所有参数')
torch.save(all_data, 'all.pth')
all_data = torch.load('all.pth')
print(all_data.keys())

版权声明:

本网仅为发布的内容提供存储空间,不对发表、转载的内容提供任何形式的保证。凡本网注明“来源:XXX网络”的作品,均转载自其它媒体,著作权归作者所有,商业转载请联系作者获得授权,非商业转载请注明出处。

我们尊重并感谢每一位作者,均已注明文章来源和作者。如因作品内容、版权或其它问题,请及时与我们联系,联系邮箱:809451989@qq.com,投稿邮箱:809451989@qq.com

热搜词