pytorch: n个GPU并行计算时模型输出的batch size等于预定义bs的n倍

我在使用DataParallel进行双GPU训练一个分类模型时,定义batch size=16,然后遇到错误:计算CrossEntropyLoss时predict与target的batch维度不一致,target的batch维度仍然是16,而模型输出的predict的batch维度竟然等于32!
按照我以前对pytorch并行化的理解,一个batch的数据应该是被平均划分到不同GPU上并行计算,然后再将结果汇总的,这样的话每个GPU上计算的batch维度大小=8,最后汇总batch size应该仍然等于定义的16。
查阅DataParallel文档,看到这么一句话:

Arbitrary positional and keyword inputs are allowed to be passed into DataParallel but some types are specially handled. tensors will be scattered on dim specified (default 0). tuple, list and dict types will be shallow copied. The other types will be shared among different threads and can be corrupted if written to in the model’s forward pass.

注意这句话:
tuple, list and dict types will be shallow copied.
然后再看我模型forward函数的输入是一个list。问题就出在这儿,两张卡并行计算时pytorch并不是将输入的list划分然后分到不同GPU中,而是直接浅拷贝后分发,这就造成同一份数据被复制n份分发到n个GPU中,而不是我预想的一份数据被平均划分为n份

再去检查模型输出的tensor,果然后16个元素的值与前16个元素的值是一样的。

因此,通过这件事,我们可以总结得到:
Pytorch中当用DataParallel或DistributedDataParallel来进行并行化计算时,模型的输入数据应该要是一个tensor的类型。

你可能感兴趣的:(pytorch: n个GPU并行计算时模型输出的batch size等于预定义bs的n倍)