pytorch中nn.DataParallel（多GPU并行计算）的原理

最新推荐文章于 2025-06-05 09:01:41 发布

道骨。

最新推荐文章于 2025-06-05 09:01:41 发布

阅读量2.8k

点赞数 3

CC 4.0 BY-SA版权

分类专栏： pytorch笔记文章标签： pytorch python

本文链接：https://blog.youkuaiyun.com/qq_41154003/article/details/115118782

pytorch笔记专栏收录该内容

1 篇文章

订阅专栏

本文详细介绍了PyTorch中的DataParallel函数，用于将模型并行化到多个GPU上进行计算。DataParallel会将batchsize的数据均匀分配到指定的GPU设备上，每个设备上都保存模型的一份副本。需要注意的是，当使用DataParallel时，如果batchsize小于GPU数量，则无法进行有效划分。函数的主要参数包括model、device_ids和output_device，其中output_device默认为第一个GPU。在实际使用中，若遇到GPU资源冲突，可通过设置CUDA_VISIBLE_DEVICES环境变量来指定可用的GPU。理解DataParallel的工作原理和配置细节对于优化深度学习模型的训练至关重要。

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

DataParallel（）函数的作用就是将一个batchsize的输入数据均分到多个GPU上分别计算。

此处注意，batchsize要大于GPU个数才能划分。

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

这个函数主要有三个参数：

module：即模型，此处注意，虽然输入数据被均分到不同gpu上，但每个gpu上都要拷贝一份模型。
device_ids：即参与训练的gpu列表，例如三块卡， device_ids = [0，1，2]
output_device：指定输出gpu，一般省略。在省略的情况下，默认为第一块卡，即索引为0的卡。此处有一个问题，输入计算是被几块卡均分的，但输出loss的计算是由这一张卡独自承担的，这就造成这张卡所承受的计算量要大于其他参与训练的卡。

此外，在使用这个函数的时候还有一个细节。

在执行DataParallel之前，会首先把其模型的参数放在device_ids[0]上，也就是参数列表里的第一张卡上。此处有坑。

例如服务器是八卡的服务器，刚好前面序号是0的卡被别人占用着，于是你只能用其他的卡来，比如你用2和3号卡，如果你直接指定device_ids=[2, 3]的话会出现模型初始化错误，类似于module没有复制到在device_ids[0]上去。那么你需要在运行train之前需要添加如下两句话指定程序可见的devices，如下：

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

当你添加这两行代码后，那么device_ids[0]默认的就是第2号卡，你的模型也会初始化在第2号卡上了，而不会占用第0号卡了。这里简单说一下设置上面两行代码后，那么对这个程序而言可见的只有2和3号卡，和其他的卡没有关系，这是物理上的号卡，逻辑上来说其实是对应0和1号卡，即device_ids[0]对应的就是第2号卡，device_ids[1]对应的就是第3号卡。