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)