nn.DataParallel：【os.environ[“CUDA_VISIBLE_DEVICES“] = “2, 3“：模型会初始化在第2号卡上，而不会占用第0号卡了】

水深无声 2023-09-27 10:58 5阅读 0赞

### 前言 ###

pytorch中的GPU操作默认是异步的，当调用一个使用GPU的函数时，这些操作会在特定设备上排队但不一定在稍后执行。这就使得pytorch可以进行并行计算。但是pytorch异步计算的效果对调用者是不可见的。

但平时我们用的更多其实是多GPU的并行计算，例如使用多个GPU训练同一个模型。Pytorch中的多GPU并行计算是数据级并行，相当于开了多个进程，每个进程自己独立运行，然后再整合在一起。

device_ids = [0, 1]
    net = torch.nn.DataParallel(net, device_ids=device_ids)

**注：多GPU计算的前提是你的计算机上得有多个GPU,在cmd上输入`nvidia-smi`来查看自己的设备上的GPU信息。**

### **nn.DataParallel详细解析** ###

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

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

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

一般我们使用torch.nn.DataParallel()这个函数来进行，接下来我将用一个例子来演示如何进行多GPU计算：

net = torch.nn.Linear(100,1)
    print(net)
    print('---------------------')
    net = torch.nn.DataParallel(net, device_ids=[0,3])
    print(net)

输出：

Linear(in_features=10, out_features=1, bias=True)
    ---------------------
    DataParallel(
      (module): Linear(in_features=10, out_features=1, bias=True)
    )

可以看到`nn.DataParallel()`包裹起来了。然后我们就可以使用这个`net`来进行训练和预测了，它将自动在第0块`GPU`和第3块`GPU`上进行并行计算，然后自动的把计算结果进行了合并。

**下面来具体讲讲`nn.DataParallel`中是怎么做的：**

首先在前向过程中，你的输入数据会被划分成多个子部分（以下称为副本）送到不同的`device`中进行计算，而你的模型`module`是在每个`device`上进行复制一份，也就是说，输入的`batch`是会被平均分到每个`device`中去，但是你的模型`module`是要拷贝到每个`devide`中去的，每个模型`module`只需要处理每个副本即可，当然你要保证你的`batch size`大于你的`gpu`个数。然后在反向传播过程中，每个副本的梯度被累加到原始模块中。概括来说就是：`DataParallel`会自动帮我们将数据切分 `load` 到相应 `GPU`，将模型复制到相应 `GPU`，进行正向传播计算梯度并汇总。

注意还有一句话，官网中是这样描述的：

The parallelized `module` must have its parameters and buffers on `device_ids[0]` before running this `DataParallel` module.

意思就是：在运行此`DataParallel`模块之前，并行化模块必须在`device_ids [0]`上具有其参数和缓冲区。在执行`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号卡。

当然你要保证上面这两行代码需要定义在下面这两行代码之前，一般放在`train.py`中`import`一些`package`之后：

device_ids = [0, 1]
    net = torch.nn.DataParallel(net, device_ids=device_ids)

那么在训练过程中，你的优化器同样可以使用nn.DataParallel，如下两行代码：

optimizer = torch.optim.SGD(net.parameters(), lr=lr)
    optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

### nn.DataParallel一些常见问题解析 ###

**1.多GPU计算减少了程序运行的时间？**

很多同学发现在进行多GPU运算时，程序花费的时间反而更多了，这其实是因为你的batch\_size太小了，因为torch.nn.DataParallel()这个函数是将每个batch的数据平均拆开分配到多个GPU上进行计算，计算完再返回来合并。这导致GPU之间的开关和通讯过程占了大部分的时间开销。

大家可以使用`watch -n 1 nvidia-smi`这个命令来查看每1s各个GPU的运行情况，如果发现每个GPU的占用率均低于50%，基本可以肯定你使用多GPU计算所花的时间要比单GPU计算花的时间更长了。

**2. 如何保存和加载多GPU网络？**

如何来保存和加载多GPU网络，它与普通网络有一点细微的不同：

net = torch.nn.Linear(10,1)  # 先构造一个网络
    net = torch.nn.DataParallel(net, device_ids=[0,3])  #包裹起来
    torch.save(net.module.state_dict(), './networks/multiGPU.h5') #保存网络
    
    # 加载网络
    new_net = torch.nn.Linear(10,1)
    new_net.load_state_dict(torch.load("./networks/multiGPU.h5"))

因为`DataParallel`实际上是一个`nn.Module`，所以我们在保存时需要多调用了一个`net.module`，模型和优化器都需要使用`net.module`来得到实际的模型和优化器。

**3. 为什么第一块卡的显存会占用的更多一些？？？**

最后一个参数`output_device`一般情况下是省略不写的，那么默认就是在`device_ids[0]`，也就是第一块卡上，也就解释了为什么第一块卡的显存会占用的比其他卡要更多一些。

进一步说也就是当你调用`nn.DataParallel`的时候，只是在你的`input`数据是并行的，但是你的`output loss`却不是这样的，每次都会在第一块`GPU`相加计算，这就造成了第一块GPU的负载远远大于剩余其他的显卡。

**4. 直接使用nn.DataParallel的时候，训练采用多卡训练，会出现一个warning？？？**

UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; 
    will instead unsqueeze and return a vector.

**首先说明一下：**

每张卡上的`loss`都是要汇总到第0张卡上求梯度，更新好以后把权重分发到其余卡。但是为什么会出现这个`warning`，这其实和`nn.DataParallel`中最后一个参数`dim`有关，其表示`tensors`被分散的维度，默认是0，`nn.DataParallel`将在`dim0`（批处理维度）中对数据进行分块，并将每个分块发送到相应的设备。单卡的没有这个`warning`，多卡的时候采用`nn.DataParallel`训练会出现这个`warning`，由于计算`loss`的时候是分别在多卡计算的，那么返回的也就是多个`loss`，你使用了多少个`gpu`，就会返回多少个`loss`。（有人建议`DataParallel`类应该有`reduce`和`size_average`参数，比如用于聚合输出的不同`loss`函数，最终返回一个向量，有多少个`gpu`，返回的向量就有几维。）

关于这个问题在pytorch官网的issues上有过讨论，下面简单摘出一些:

[https://github.com/pytorch/pytorch/issues/9811github.com/pytorch/pytorch/issues/9811][https_github.com_pytorch_pytorch_issues_9811github.com_pytorch_pytorch_issues_9811]

前期探讨中，有人提出求`loss`平均的方式会在不同数量的`gpu`上训练会以微妙的方式影响结果。模块返回该`batch`中所有损失的平均值，如果在4个`gpu`上运行，将返回4个平均值的向量。然后取这个向量的平均值。但是，如果在3个`GPU`或单个`GPU`上运行，这将不是同一个数字，因为每个`GPU`处理的`batch size`不同！举个简单的例子（就直接摘原文出来）：

> A batch of 3 would be calculated on a single GPU and results would be \[0.3, 0.2, 0.8\] and model that returns the loss would return 0.43.  
> If cast to DataParallel, and calculated on 2 GPUs, \[GPU1 - batch 0,1\], \[GPU2 - batch 2\] - return values would be \[0.25, 0.8\] (0.25 is average between 0.2 and 0.3)- taking the average loss of \[0.25, 0.8\] is now 0.525!  
> Calculating on 3 GPUs, one gets \[0.3, 0.2, 0.8\] as results and average is back to 0.43!

似乎一看，这么求平均loss确实有不合理的地方。那么有什么好的解决办法呢，可以使用`size_average=False`，`reduce=True`作为参数。每个`GPU`上的损失将相加，但不除以`GPU`上的批大小。然后将所有平行损耗相加，除以整批的大小，那么不管几块`GPU`最终得到的平均`loss`都是一样的。

那`pytorch`贡献者也实现了这个`loss`求平均的功能，即通过`gather`的方式来求loss平均：

[https://github.com/pytorch/pytorch/pull/7973/commits/c285b3626a7a4dcbbddfba1a6b217a64a3f3f3begithub.com/pytorch/pytorch/pull/7973/commits/c285b3626a7a4dcbbddfba1a6b217a64a3f3f3be][https_github.com_pytorch_pytorch_pull_7973_commits_c285b3626a7a4dcbbddfba1a6b217a64a3f3f3begithub.com_pytorch_pytorch_pull_7973_commits_c285b3626a7a4dcbbddfba1a6b217a64a3f3f3be]

如果它们在一个有2个`GPU`的系统上运行，`DP`将采用多`GPU`路径，调用`gather`并返回一个向量。如果运行时有1个`GPU`可见，`DP`将采用顺序路径，完全忽略`gather`，因为这是不必要的，并返回一个标量。

[\[1\]][1][\[2\]][2]

### 参考 ###

1.  [^][Link 1]Pytorch多GPU计算之torch.nn.DataParallel() [Pytorch多GPU计算之torch.nn.DataParallel()\_miguemath的博客-CSDN博客][Pytorch_GPU_torch.nn.DataParallel_miguemath_-CSDN]
2.  [^][Link 2]Pytorch的nn.DataParallel [Pytorch的nn.DataParallel - 知乎][Pytorch_nn.DataParallel -]

[https_github.com_pytorch_pytorch_issues_9811github.com_pytorch_pytorch_issues_9811]: https://link.zhihu.com/?target=https%3A//github.com/pytorch/pytorch/issues/9811
[https_github.com_pytorch_pytorch_pull_7973_commits_c285b3626a7a4dcbbddfba1a6b217a64a3f3f3begithub.com_pytorch_pytorch_pull_7973_commits_c285b3626a7a4dcbbddfba1a6b217a64a3f3f3be]: https://link.zhihu.com/?target=https%3A//github.com/pytorch/pytorch/pull/7973/commits/c285b3626a7a4dcbbddfba1a6b217a64a3f3f3be
[1]: https://zhuanlan.zhihu.com/p/393857045#ref_1
[2]: https://zhuanlan.zhihu.com/p/393857045#ref_2
[Link 1]: https://zhuanlan.zhihu.com/p/393857045#ref_1_0
[Pytorch_GPU_torch.nn.DataParallel_miguemath_-CSDN]: https://blog.csdn.net/wangkaidehao/article/details/104411682
[Link 2]: https://zhuanlan.zhihu.com/p/393857045#ref_2_0
[Pytorch_nn.DataParallel -]: https://zhuanlan.zhihu.com/p/102697821