Pytorch:多块GPU调用细节问题及Pytorch的nn.DataParallel解释

2022-10-06 11:09:44

我们用实验室带有多块卡的GPU服务器,当我们在上面跑程序的时候,当迭代次数或者epoch足够大的时候,我们通常会使用nn.DataParallel函数加入以下代码段来用多个GPU来加速训练。

device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

然鹅,也常常会由于正在运行程序的0卡显存不够放不下我们新来的模型初始化数据而导致显存OOM的问题不断,使得模型无法训练,同时报出警告:

UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars;will instead unsqueeze and return a vector.

下面就GPU存放数据进行整理说明,我们在ternimal下执行语句下列语句后,

watch -n 1 nvidia-smi

我们会发现确实会使用多个GPU来并行训练。但是同时我们会发现其实第一块卡的显存会占用的更多一些。那是因为官方有定义,详细可以参看DataParallel的函数定义。

CLASS torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

在前向过程中,我们的输入数据会被划分成多个子部分(可称为副本)送到不同的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 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号卡。

上面这两行代码需要定义在

device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

这两行代码之前,一般放在train.py中import一些package之后。

那么在训练过程中,优化器同样可以使用nn.DataParallel,如下两行代码:

optimizer = torch.optim.SGD(net.parameters(), lr=lr)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

那么使用nn.DataParallel后,事实上DataParallel也是一个Pytorch的nn.Module,那么模型和优化器都需要使用.module来得到实际的模型和优化器,如下:

保存模型:
torch.save(net.module.state_dict(), path)
加载模型:
net=nn.DataParallel(Resnet18())
net.load_state_dict(torch.load(path))
net=net.module
优化器使用:
optimizer.step() --> optimizer.module.step()

还有一个问题就是,如果直接使用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上讨论过:DataParallel does not work with tensors of dimension 0

  • 作者:yunxiaoMr
  • 原文链接:https://blog.csdn.net/weixin_41297324/article/details/113361394
    更新时间:2022-10-06 11:09:44