Pytorch单机多卡训练时的坑

lzl2040

已于 2022-10-17 17:07:19 修改

阅读量2.1k

点赞数

CC 4.0 BY-SA版权

分类专栏： python 计算机视觉我的笔记文章标签： pytorch 深度学习 python

于 2022-10-17 17:05:52 首次发布

本文链接：https://blog.youkuaiyun.com/qq_41234663/article/details/127367744

我的笔记同时被 3 个专栏收录

46 篇文章

订阅专栏

计算机视觉

10 篇文章

订阅专栏

python

5 篇文章

订阅专栏

博客讲述了在Pytorch中进行单机多GPU训练时遇到的问题，即模型在多GPU训练后无法在单GPU上正确验证，原因是训练时模型参数带有module。解决方案包括在保存模型时针对多GPU训练的情况使用`.module.state_dict()`，以及在验证时使用`nn.DataParallel`或`nn.parallel.DistributedDataParallel`加载和验证模型。此外，提供了相应的代码片段以供参考。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Pytorch单机多卡训练时的坑

前言

昨天训练了一天的模型，今天来验证发现出问题了，研究发现是因为我训练的时候是多GPU，但验证的时候是单GPU，导致模型不一致。

遇到的问题

验证的时候出现了如下的错误：

Missing key(s) in state_dict: "key_encoder.conv1.weight" .....

这个的原因就是因为训练的时候是多GPU，但验证的时候使用的是单GPU

解决办法

（1）可以在训练的时候在保存模型的地方做些修改：
对于单GPU训练的，模型保存方式为：

torch.save(self.RDE_VOS.state_dict(), model_path)

对于多GPU训练的，模型保存方式为：

torch.save(self.RDE_VOS.module.state_dict(), model_path)

两者区别在于一个module,因为进行多GPU训练时会在模型的参数上加一个module,这点就很坑，如果是初次使用分布式训练的同学会很容易进坑

（2）如果已经训练好了的，那我们也有办法，就是还是使用多卡验证

对于使用nn.DataParallel训练的

# 将模型使用数据并行方式加载
net = nn.DataParallel(net, device_ids=[0, 1, 2, 3])
net.load_dict(torch.load(args.weights))
net.eval()

对于使用nn.parallel.DistributedDataParallel训练的

distributed.init_process_group(backend="nccl")
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
prop_model = nn.parallel.DistributedDataParallel(RDE_VOS().cuda(),device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False)
map_location = 'cuda:%d' % self.local_rank
path = "你的模型地址"
src_dict = torch.load(path, map_location={'cuda:0': map_location})
prop_model.load_state_dict(src_dict, strict=False)
prop_model.eval()

然后在命令行输入语句运行即可。

python -m  torch.distributed.launch --master_port 9843  --nproc_per_node=2 eval_davis_2016.py