Pytorch 的 nn.DataParallel 详细解析 - 极市社区
source link: https://bbs.cvmart.net/articles/5250
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
作者丨初识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)
:
这个函数主要有三个参数:
module
:即模型,此处注意,虽然输入数据被均分到不同gpu上,但每个gpu上都要拷贝一份模型。device_ids
:即参与训练的gpu列表,例如三块卡, device_ids = [0,1,2]。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.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
前期探讨中,有人提出求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平均:
如果它们在一个有2个GPU
的系统上运行,DP
将采用多GPU
路径,调用gather
并返回一个向量。如果运行时有1个GPU
可见,DP
将采用顺序路径,完全忽略gather
,因为这是不必要的,并返回一个标量。
- 2
- 1
- 2066
Recommend
About Joyk
Aggregate valuable and interesting links.
Joyk means Joy of geeK