Pytorch多GPU并行Bug收集（长期）：KeyError: Caught KeyError in replica *[device_id] on device *[device_id].

在单卡训练顺利的前提下，修改为多卡训练，可谓bug多多

今天用pytorch 多GPU并行训练时，在最后一个step的时候报错了，

KeyError: Caught KeyError in replica 5 on device 5.

如图所示

我们可以看到，代码处应该是没有问题的，经过我的计算，利用train的总量/batch_size-->16100/24=670---余20，无法整除，这下有点眉目了。正好出现在最后一个step上。

我这次试验用的是6个GPU跑的试验，使用torch.nn.DataParallel来实现的并行，按照源码的并行逻辑，是会根据num_parallel=batch_size/gpu_nums来决定均匀分配给每块GPU上的样本数量，那么我这里num_parallel=4.

到这里有没有突然发现数字上有了联系。没错，20/4=5。正好最后一个step的时候，前5张GPU把数据分完了，最后一张GPU上没有分到数据。

经过合理计算，修正batch_size和gpu个数，使得最后一个step时每张GPU上都能分到数据。

问题解决完毕。