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训练的模型
      </div>