pytorch持久化

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

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

torch.savetorch.load

python 复制代码
import torch

# 1.变量的保存与加载
a = torch.Tensor(3,4)
print(a.get_device())  # -1,代表CPU
if torch.cuda.is_available():
    a = a.cuda()   # 把a转换为GPU0上的tensor
    torch.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())
相关推荐
学术小八36 分钟前
2025年人工智能、虚拟现实与交互设计国际学术会议
人工智能·交互·vr
仗剑_走天涯2 小时前
基于pytorch.nn模块实现线性模型
人工智能·pytorch·python·深度学习
cnbestec3 小时前
协作机器人UR7e与UR12e:轻量化设计与高负载能力助力“小而美”智造升级
人工智能·机器人·协作机器人·ur协作机器人·ur7e·ur12e
zskj_zhyl3 小时前
毫米波雷达守护银发安全:七彩喜跌倒检测仪重构居家养老防线
人工智能·安全·重构
gaosushexiangji4 小时前
利用sCMOS科学相机测量激光散射强度
大数据·人工智能·数码相机·计算机视觉
ai小鬼头5 小时前
AIStarter新版重磅来袭!永久订阅限时福利抢先看
人工智能·开源·github
说私域5 小时前
从品牌附庸到自我表达:定制开发开源AI智能名片S2B2C商城小程序赋能下的营销变革
人工智能·小程序
飞哥数智坊6 小时前
新版定价不够用,Cursor如何退回旧版定价
人工智能·cursor
12点一刻6 小时前
搭建自动化工作流:探寻解放双手的有效方案(2)
运维·人工智能·自动化·deepseek
未来之窗软件服务6 小时前
东方仙盟AI数据中间件使用教程:开启数据交互与自动化应用新时代——仙盟创梦IDE
运维·人工智能·自动化·仙盟创梦ide·东方仙盟·阿雪技术观