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())
相关推荐
MARS_AI_5 分钟前
大语言模型驱动智能语音应答:技术演进与架构革新
人工智能·语言模型·自然语言处理·架构·信息与通信
程序员小灰10 分钟前
AI独角兽团队Manus裁员80人,剩下40人迁至新加坡总部!
人工智能·aigc·agent
新智元21 分钟前
OpenAI去年挖的坑填上了!奖励模型首现Scaling Law,1.8B给70B巨兽上了一课
人工智能·openai
简婷1870199877531 分钟前
源网荷储 + 零碳园区:一场关于能源与未来的双向奔赴
大数据·人工智能·能源
新智元35 分钟前
Grok 4作战图刷爆全网,80%华人横扫硅谷!清华上交校友领衔,95后站C位
人工智能·openai
小宋00135 分钟前
使用LLaMA-Factory微调Qwen2.5-VL-3B 的目标检测任务-数据集格式转换(voc 转 ShareGPT)
人工智能·目标检测·计算机视觉
小哥谈1 小时前
论文解析篇 | YOLOv12:以注意力机制为核心的实时目标检测算法
人工智能·深度学习·yolo·目标检测·机器学习·计算机视觉
mit6.8241 小时前
[Meetily后端框架] AI摘要结构化 | `SummaryResponse`模型 | Pydantic库 | vs marshmallow库
c++·人工智能·后端
阿里云大数据AI技术1 小时前
Post-Training on PAI (3): 自研高性能强化学习框架PAI-ChatLearn
人工智能·开源·强化学习
二二孚日1 小时前
自用华为ICT云赛道AI第三章知识点-MindSpore特性、MindSpore开发组件
人工智能·华为