pytorch多卡训练显卡利用率问题

多卡训练显卡利用率问题

最近我继承了前同事的深度学习代码，使用pytorch写的。在nvidia-smi查看显卡利用率的时候发现显卡是经常出现除了第一张显卡外，其他7张显卡的利用率为0的情况。

同时查看了CPU利用率，发现大多的核也是空闲的：

阅读代码后先后试了：调整dataloader的num_workers的数量（之前为默认值）、把数据先加载到内存的方法均无法得到明显改善。

然后我debug各个阶段的耗时，发现在inference时耗时0.1秒多，但是在计算loss和回传loss时总共用时接近6秒（图丢了）。所以我推测代码把所有卡的loss都对到第一张卡里去计算了，并且在代码中得到了验证，前同事代码中训练函数的三大件儿是这么写的：

net_output = model(inputs)
optimizer.zero_grad()
loss_all = compute_loss(net_output, targets_semantic[:, 0])
loss = loss_all['total_loss']
loss.backward()
optimizer.step()

可以看到，他是把模型的inference和计算loss是分开写了。这也是大多教程中的一种写法，单卡训练的时候是没有任何问题的，并且还会使得代码结构更加清晰。但是在多卡训练时这就导致程序在inference完成之后把八张卡的预测结果丢到了第一张卡里去计算loss并运行loss.backward 这就导致这两块的时间消耗巨大。

随后我将代码进行调整，把计算loss的功能放到的model类中的forward里，使得在八张卡中分别计算loss。这样可以使显卡利用率达到最大。修改后的训练函数三大件儿这么写的：

output = model(inputs, targets_semantic[:, 0])
net_output = output['typ_pred']
loss = output['loss'] # 返回的是一个List
loss = loss.mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()

即，八张卡分别计算loss，再将loss求均值，最后再在第一张卡中做回传。

这种做法之后显卡利用率得到大幅度提升，在我的代码上大改提升了3-4倍的训练速度（不贴图了懒得做实验）。

但是仍然存在问题是：Loss.backward是否也可以在8张卡上分别做回传？这样就不会出现第一张卡在运行loss.backward的时候其余七张卡在短暂空闲的问题了。但是在多卡上分别Loss.backward的话会使得各显卡上模型丧失通讯，对最终的模型影响有多大？这是后续需要优化的问题。