本文正在参与「金石方案 . 分割6万现金大奖」

前言

有一期的歹意文件检测模型练习好了,因此需求进行测验,关于歹意文件检测的内容,可以回看博主之前写的博文:

  • 【AI】浅析歹意文件静态检测及部分问题处理思路
  • 【AI】歹意文件静态检测模型检验及小结

因为样本在某台机子上,又恰逢有其他模型在练习,因此 GPU 资源被占满了,不过测验这个模型的话,CPU 也绰绰有余了,当我准备运用 CPU 练习时,却遇到了问题;

剖析

1、model.to(device) 不会影响 torch.load()

我一开始认为只需运用 model.to 就算是运用上 CPU 了;

device = torch.device("cpu")
model = ...
model = model.to(device)
model_savedir_ = ''
if os.path.exists(model_savedir_):
    print("model load.")
    state_dict = torch.load(model_savedir_)
    model.load_state_dict(state_dict)

事实证明,我想的太简略了…

【问题解决】解决如何在 CPU 上加载多 GPU 训练的模型

RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

这个问题很显而易见,便是 GPU 的内存溢出了,但是按我的思路,用的应该是 CPU 啊,所以我怀疑是 torch.load() 这个函数出了问题,查询了一番资料后,发现是要这样运用的 state_dict = torch.load(model_savedir_, map_location=device)


2、GPU 与 CPU 练习时参数名不一致

当我认为功德圆满,点击运转之时,不料,又报错了:

【问题解决】解决如何在 CPU 上加载多 GPU 训练的模型

RuntimeError: Error(s) in loading state_dict for ..model..:
	Missing key(s) in state_dict: "fc.weight", "fc.bias", "features.0.0.weight", "features.0.1.weight", "features.0.1.bias", "features.0.1.running_mean", "features.0.1.running_var", "features.1.conv.0.weight", "features.1.conv.1.weight", "features.1.conv.1.bias", "features.1.conv.1.running_mean", "features.1.conv.1.running_var", "features.1.conv.3.weight", "features.1.conv.4.weight", "features.1.conv.4.bias", "features.1.conv.4.running_mean", "features.1.conv.4.running_var", "features.1.conv.5.fc.0.weight", ...

依据了解,便是说找不到参数,因此,我将字典部分内容打印了一下:

for k, v in state_dict.items():
    print(k, v)
    break

【问题解决】解决如何在 CPU 上加载多 GPU 训练的模型

发现问题了,在多 GPU 上练习的模型,保存时会在参数名前多加了一个 module. 前缀,因此在用 CPU 进行加载时,需求把这个前缀去掉:

if os.path.exists(model_savedir_):
    print("model load.")
    state_dict = torch.load(model_savedir_, map_location=device)
    from collections import OrderedDict
    state_dict_new = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:]  # 去掉 `module.`
        state_dict_new[name] = v
    model.load_state_dict(state_dict_new)

这样就可以在 CPU 上加载多 GPU 练习的模型了!

后记

以上便是 【问题处理】处理如何在 CPU 上加载多 GPU 练习的模型 的全部内容了,期望对大家有所协助!

上篇精讲:【问题处理】处理 Docker 二次重启 MySQL 8 遇到的一些问题

我是,等待你的关注;

创作不易,请多多支撑;

系列专栏:问题处理 AI