文章目录
- 1 问题描述
- 2 模型保存方式
- 3 单块GPU上加载模型
- 4 CPU上加载模型
- 5 总结
1 问题描述
PyTorch提供了非常便捷的多GPU网络训练方法：DataParallel和DistributedDataParallel。在涉及到一些复杂模型时，基本都是采用多个GPU并行训练并保存模型。但在推理阶段往往只采用单个GPU或者CPU运行。这时怎么将多GPU环境下保存的模型权重加载到单GPU/CPU运行环境下的模型上成了一个关键的问题。
如果忽视环境问题直接加载往往会出现两类问题：
1 出现错误：IndexError: list index out of range
出现这个错误的原因是：现有模型参数是在多GPU上获得并保存的，因此在读入时默认会保存至对应的GPU上，但是目前推理环境中只有一块GPU，所以导致那些本来在其它GPU上的参数找不到自己应该去的GPU编号，出现了一个溢出错误，本质是GPU编号溢出。
2 出现错误：Missing key(s) in state_dict:
出现这个错误的原因是：由于模型训练和推理的环境不同，导致一些参数丢失，因此报错。目前在网上的一些解决策略是忽视这些丢失的参数，例如使用命令：model.load_state_dict(torch.load('model.pth'), strict=False)
来成功导入模型。这条命令可以让程序不报错并看似成功的导入模型参数。但实际上这条命令的含义是在导入模型参数时通过设置 strict=False 来忽略丢失的参数，也就是说那些丢失参数地方的模型权重初仍为初始化随机状态，等同于没有进行训练和学习，何谈推理与验证！！！
2 模型保存方式
不论是用哪种方式进行推理，在训练的时候要保证程序保存模型的方式是这样的：
torch.save(model.state_dict(), "model.pth")
3 单块GPU上加载模型
将多GPU训练的权重文件加载到单GPU上：
# 1 加载模型
model = Model()
# 2 指定运行设备，这里为单块GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 3 将模型用DataParallel方法封装一次
model = torch.nn.DataParallel(model)
# 4 将模型读入到GPU设备上
model = model_E2E.to(device)
# 5 加载权重文件
model.load_state_dict(torch.load(weight_path, map_location=device))
通过上面的程序就可以实现将多块GPU上训练得到的权重文件加载到单块GPU环境下的模型中。这里有两点需要注意：
-
在多GPU训练时，模型使用了
DataParallel或
DistributedDataParallel方法，这两种并行化工具会修改模型的结构，将模型封装在一个新的模块中，通常名为：module。因此在权重文件中保存的模型是经过DataParallel封装后的结构。为了能够载入全部参数，需要通过步骤3使推理模型与原始多GPU训练模型在结构上保持一致。 -
在步骤5加载模型参数时使用了
map_location参数。这个参数会告诉 PyTorch在加载模型时应该将张量放置在哪个设备上。设置map_location=device，那么无论模型原来是在哪个设备上训练的，现在都将放置在指定的设备device='cuda:0'上。
4 CPU上加载模型
在CPU上加载模型：
from collections import OrderedDict
# 1 加载模型
model = Model()
# 2 指定设备CPU
device = "cpu"
# 3 读取权重文件
state_dict = torch.load(weight_path, map_location=device)
# 4 剥除权重文件中的module层
if next(iter(state_dict)).startswith("module."):
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
state_dict = new_state_dict
# 5 加载权重文件
model.load_state_dict(state_dict)
# 6 将模型载入到CPU
model = model.to(device)
在CPU上加载模型的逻辑和GPU的差不多，核心都是因为原权重文件中的模型被封装成了module.Model，所以需要将这层外壳去掉，最后再进行读取并将模型加载到CPU上。
5 总结
在深度学习任务中训练与推理环境存在差异的情况十分常见 ，有差异的环境下实现网络权重文件的正确读取十分重要。实际操作中一定要确保正确的权重文件被读入，这是进行推理最基本的前提！最好在推理前做一些对比实验（例如：选取一部分数据，分别套用已有的程序进行训练和推理，对比二者的效果）来确保已经读入到正确的权重。
