模型的保存与加载与多gpu的模型保存和加载

模型保存与加载

模型的保存与加载方式

• 模型保存有两种形式，一种是保存模型的state_dict()，只是保存模型的参数。那么加载时需要先创建一个模型的实例model，之后通过torch.load()将保存的模型参数加载进来，得到dict，再通过model.load_state_dict(dict)将模型的参数更新。
• 另一种是将整个模型保存下来，之后加载的时候只需要通过torch.load()将模型加载，即可返回一个加载好的模型。

一、方式一：保存和加载模型参数

1、模型保存

import torch
import torchvision.models as models
model = self.model.module if isinstance(self.model, torch.nn.DataParallel) else self.model
model_path = "model_weights.pth"
torch.save(model.state_dict(), model_path)

2、模型加载：当加载模型权重时，我们需要首先实例化模型类，因为类定义了网络的结构。

想要加载模型参数，我们需要创建一个和原模型一样的实例，然后通过load_state_dict()方法来加载模型参数。
由上可知当你用state_dict()保存模型参数后，用torch.load(“xx.pth”）得到的是个dict

model = models.vgg16() # we do not specify pretrained=True, i.e. do not load default weights
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

二、方式二：保存和加载模型参数与结构

当加载模型权重时，需要首先实例化模型类，因为类定义了网络的结构。
我们可能希望将这个类的结构与模型保存在一起。这样的话，我们可以将model而不是model.state_dict()作为参数。

torch.save(model, 'model.pth')

这样，我们加载模型的时候就不用再新建一个实例了。加载方式如下所示

model = torch.load('model.pth')

• 这种方式在网络比较大的时候可能比较慢，因为相较于上面的方式多存储了网络的结构

加载预训练模型中的部分参数，并固定该部分参数

一、加载预训练模型里我们初始化模型也需要的参数（可适用于断点之后，要加载那个模型再训练，如CMT的trainer.load_model里写的）

注意：需要重新训练的层的名字要和之前的不同。

model=resnet()#自己构建的模型，以resnet为例

model_dict = model.state_dict()  #加载实例化模型的初始参数

pretrained_dict = torch.load('xxx.pth')

#之前预训练模型的参数，挑选出实例化模型里面有的参数，放进pretrained_dict 
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}  

#然后把预训练时训练过的一些参数，更新到初始化参数model_dict字典里
model_dict.update(pretrained_dict)

#最后把已经更新过的，加载了部分预训练参数的参数字典用load_state_dict写入model,将模型参数更新
model.load_state_dict(model_dict)

二、如果想固定部分参数

#k是可训练参数的名字，v是包含可训练参数的一个实体

#可以先print（k）,找到自己想进行调整的层，并将该层的名字加入到if语句中：
for k,v in model.named_parameters():

    if k!='xxx.weight' and k!='xxx.bias' :

        v.requires_grad=False#固定参数

训练部分参数

#将要训练的参数放入优化器

optimizer2=torch.optim.Adam(params=[model.xxx.weight，model.xxx.bias],lr=learning_rate,betas=(0.9,0.999),weight_decay=1e-5)

多GPU

转载多gpu模型保存与加载

一、想用多GPU同时训练

只需下面一步简单的操作：使用模块DataParallel，你就可以很容易地在多个gpu上运行你的操作:

if torch.cuda.device_count() > 1:
     print("Let's use", torch.cuda.device_count(), "GPUs!")
     model = torch.nn.DataParallel(model)

二、多GPU环境下对模型的保存

用model.module模块去保存模型。
要先判断下是否为数据并行的形式，多GPU情况下，需要调用model.module模块，再去保存模型参数state_dict。我理解model.module就是将多卡上的参数以及网络机构通过某种机制将其汇总成：只有一个网络结构，只有一套参数的模型结构。

#万能的保存方法，如果你的预测函数不会依赖你的模型类定义。
if isinstance(model, torch.nn.DataParallel):
   torch.save(model.module.state_dict(), config.save_path)
#依赖你的模型定义，这么保存只保存了你模型的参数，模型的结构没有保存，所以尽量用上面的保存方法。
if isinstance(model, torch.nn.DataParallel):
   torch.save(model.state_dict(), config.save_path)
 

三、多GPU环境下对模型的加载去预测

在 pytorch 多GPU训练下，存储 整个模型 （ 而不是model.state_dict() ）后再调用模型可能会遇到下面的情况：

AttributeError: ‘DataParallel’ object has no attribute ‘xxxx’

解决的方法是：

model = torch.load('path/to/model')
if isinstance(model,torch.nn.DataParallel):
		model = model.module
 
#下面就可以正常使用了
model.eval()

三、多GPU环境下LOSS的计算

先来看下loss的反向传播代码：

for step, batch in enumerate(train_dataloader):
	batch[0] = torch.LongTensor(batch[0]).to(config.device)
    batch[1] = torch.LongTensor(batch[1]).to(config.device)
    loss, logits= model(batch[0],labels=batch[1])   # Forward pass
    loss_avg.backward()  # Backward pass

报错如下，意思就是梯度只能为标量（即一个数）输出隐式地创建。

RuntimeError: grad can be implicitly created only for scalar outputs

这里面4gpu打印输出的是一个list，即每一个gpu上都会有一个loss返回：

---loss_avg : tensor([0.5866, 0.6941, 0.6176, 0.6226], device='cuda:0',
       grad_fn=<DivBackward0>)
---loss_avg : tensor([0.5866, 0.6941, 0.6176, 0.6226], device='cuda:0',
       grad_fn=<DivBackward0>)

平时我们单卡的loss形式：

---loss_avg : tensor(0.4542, device='cuda:0', grad_fn=<DivBackward0>)
---loss_avg : tensor(0.3607, device='cuda:0', grad_fn=<DivBackward0>)

解决方法：

将多卡得到的loss进行mean，求平均：

loss_avg.backward()  -----> loss_avg.mean().backward()