• 问答
  • 技术
  • 实践
  • 资源
Pytorch 的 nn.DataParallel 详细解析精选
技术讨论

作者丨初识CV@知乎(已授权)
来源丨https://zhuanlan.zhihu.com/p/393857045
编辑丨极市平台

前言

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](https://link.zhihu.com/?target=https%3A//pytorch.org/docs/stable/nn.html%3Fhighlight%3Dtorch%2520nn%2520datapa%23torch.nn.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.pyimport一些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类应该有reducesize_average参数,比如用于聚合输出的不同loss函数,最终返回一个向量,有多少个gpu,返回的向量就有几维。)

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

https://github.com/pytorch/pytorch/issues/9811​github.com

前期探讨中,有人提出求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=Falsereduce=True作为参数。每个GPU上的损失将相加,但不除以GPU上的批大小。然后将所有平行损耗相加,除以整批的大小,那么不管几块GPU最终得到的平均loss都是一样的。

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

https://github.com/pytorch/pytorch/pull/7973/commits/c285b3626a7a4dcbbddfba1a6b217a64a3f3f3be​github.com

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

参考链接

  1. Pytorch多GPU计算之torch.nn.DataParallel
  2. Pytorch的nn.DataParallel
  • 3
  • 1
  • 2281
收藏
暂无评论
chengzi

阿里巴巴

  • 62

    关注
  • 205

    获赞
  • 15

    精选文章
近期动态
  • 目标跟踪
文章专栏
  • 目标跟踪分享