Pytorch 1.0 多gpu训练尝试

最新推荐文章于 2025-06-15 11:24:16 发布

原创最新推荐文章于 2025-06-15 11:24:16 发布 · 4.7k 阅读

5 ·

CC 4.0 BY-SA版权

pytorch 专栏收录该内容

5 篇文章

订阅专栏

本文介绍了在PyTorch 1.0中进行多GPU训练时遇到的问题，如负载不均衡和训练速度慢。通过探讨官方API和张航大神的解决方案，总结了当前并行训练的局限性，并提出了解决负载不均衡的可能途径。

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

单GPU

开始用过全局指定os.environ["CUDA_VISIBLE_DEVICES"] = "0"的方法，但是运行会爆奇怪的cudnn错误，后面改成下面写法运行通过

my_device="cuda:0" #这里的0可以根据需要改为指定的gpu id
#对于模型
model=model.to(device=my_device)
#对于数据
input_data=Variable(input_data.to(device=my_device))

多GPU

实验室目前只有三块2080ti可用，所以使用三块做测试
开始使用nn.DataParallel发现，出现负载不均衡问题，且loss优化缓慢，并且运行速度竟然比单卡训练还慢

device_ids =[0,1,2]#这里设置好可用的gpu id
# 对于模型
model=model.to(device=device_ids[0])
model=nn.DataParallel(model,device_ids=device_ids)
# 对于数据
input_data=Variable(input_data.to(device=device_ids[0]))
# 对于优化器
optimizer=nn.DataParallel(optimizer,device_ids=device_ids) #优化器
optimizer.module.step()

后采用nn.parallel.DistributedDataParallel,虽然这个api原来设计给分布式使用，但是对于单机多gpu依然有效。

#初始化使用nccl后端
torch.distributed.init_process_group(backend="nccl")
#模型并行化
model=torch.nn.parallel.DistributedDataParallel(model)

单机运行时使用下面方法启动

python -m torch.distributed.launch main.py

总结：目前pytorch提供的并行方案都存在一定的负载不均衡问题，正向传播并行，但是计算loss和反向传播时交给主GPU，所以当模型较小和数据量较小时，对比单GPU运行毫无优势。解决负载均衡问题可以参考官方nn.parallel.DistributedDataParallel的文档，自行编写并行的loss方法。
或参考张航大神给出的解决方案
https://hangzhang.org/PyTorch-Encoding/notes/compile.html
但目前由于支持最新pytorch 1.0版本，并且还蛮多bug的，表示我还并没有运行成功。暂时先mark。