Pytorch 模型的加载与保存

2024年4月19日 210次阅读 来源: pytorch

pytorch的模型和参数是分开的，可以分别保存或加载模型和参数。

1、直接保存模型

# 保存模型
torch.save(model, 'model.pth')
# 加载模型
model = torch.load('model.pth')

2、分别加载模型的结构和参数

# 保存模型参数
torch.save(model.state_dict(), 'model.pth')
# 加载模型参数
model.load_state_dict(torch.load('model.pth')

CPU模型加载GPU参数

model.load_state_dict(torch.load('model.pth', map_location='cpu'))

通过DataParalle使用多GPU

model=DataParalle(model)
#保存参数
torch.save(model.module.state_dict(), 'model.pth')

   

自己习惯用的代码段


# 判断gpu是否可用
use_cuda = torch.cuda.is_available()
# 是否使用多gpu
use_multi_gpu = True
# 默认加载的cpu的参数
model.load_state_dict(torch.load('model.pth')

if use_cuda:
    model = model.cuda()
if use_multi_gpu:
    model = DataParalle(model)

# 保存模型参数（一般保存cpu的参数比较好）
if use_multi_gpu:
    torch.save(model.cpu().module.state_dict(), 'model.pth')
else:
    torch.save(model.cpu().state_dict(), 'model.pth')

     
3、pytorch预训练模型
加载预训练模型和参数
resnet18 = models.resnet(pretrained=True)

     
只加载模型，不加载预训练参数
# 加载模型
resnet18 = models.resnet18(pretrained=False)
# 加载预先下载好的预训练模型参数
resnet18.load_state_dict(torch.load('resnet18-5c106cde.pth'))

     
加载部分预训练模型
resnet152 = models.resnet152(pretrained=True)
pretrained_dict = resnet152.state_dict()
"""加载torchvision中的预训练模型和参数后通过state_dict()方法提取参数 也可以直接从官方model_zoo下载： pretrained_dict = model_zoo.load_url(model_urls['resnet152'])"""
model_dict = model.state_dict()
# 将pretrained_dict里不属于model_dict的键剔除掉
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 更新现有的model_dict
model_dict.update(pretrained_dict)
# 加载我们真正需要的state_dict
model.load_state_dict(model_dict)

     
【参考资料】
 1、PyTorch学习：加载模型和参数
 2、PyTorch使用cpu调用gpu训练的模型
    原文作者：pytorch
    原文地址: https://www.cnblogs.com/leebxo/p/10920134.html
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。
点赞