Pytorch多GPU训练踩坑记录2

问题介绍

使用nn.DataParallel进行多GPU训练时，对模型进行传参，有时会出现报错“RuntimeError: chunk expects at least a 1-dimensional tensor”。

问题分析

nn.DataParallel的作用是将模型和数据分配到各个GPU上，让其在各自的GPU上训练，首先检查batchsize是否是GPU数量的整数倍，以及dataloader是否设置了drop_last是否是true,如果没有drop_last，那么最后一个batchsize剩下的数量可能不是GPU数量的整数倍，那么也会出现此错。如果以上情况都不是报错原因，还有一个原因可能是：对模型进行传参时，有可能传递了标量。

解决方案

nn.DataParallel没法将标量分散然后分配给各个GPU，因此，多GPU训练时，最好避免传参为标量，一定要传递这个此参数的话，最好在nn.DataParallel之前传参，或者最好把参数改为张量。