【Pytorch】多GPU训练

在具体使用pytorch框架进行训练的时候，发现实验室的服务器是多GPU服务器，因此需要在训练过程中，将网络参数都放入多GPU中进行训练。

   正文开始：

   涉及的代码为torch.nn.DataParallel，而且官方推荐使用nn.DataParallel而不是使用multiprocessing。官方代码文档如下：nn.DataParallel   教程文档如下：tutorial

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

该函数实现了在module级别上的数据并行使用，注意batch size要大于GPU的数量。

参数 ： module：需要多GPU训练的网络模型

device_ids： GPU的编号（默认全部GPU）

output_device：（默认是device_ids[0])

dim：tensors被分散的维度，默认是0

在代码文档中使用方法为：

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
out = net(input)

具体实际操作中，加入其他代码会更好一点。

重点：

自己的代码及解释如下：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

定义device，其中需要注意的是“cuda:0”代表起始的device_id为0，如果直接是“cuda”，同样默认是从0开始。可以根据实际需要修改起始位置，如“cuda:1”。

model = Model()
if torch.cuda.device_count() > 1:
  model = nn.DataParallel(model，device_ids=[0,1,2])
 
model.to(device)

这里注意，如果是单GPU，直接model.to(device)就可以在单个GPU上训练，但如果是多个GPU就需要用到nn.DataParallel函数，然后在进行一次to(device)。

需要注意：device_ids的起始编号要与之前定义的device中的“cuda:0”相一致，不然会报错。

如果不定义device_ids，如model = nn.DataParallel(model)，默认使用全部GPU。定义了device_ids就可以使用指定的GPU，但一定要注意与一开始定义device对应。

通过以上代码，就可以实现网络的多GPU训练。

以下是训练过程中遇到其他的bug：

1、在实际训练过程中，如果要用到网络中的子模块，需要注意：

在单GPU中，可以使用以下代码

model = Net()
out = model.fc(input)

但是在DataParallel中，需要修改为如下：

model = models.resnet34(True)
model.fc = nn.Linear(model.fc.in_features, 2)
model.cuda()
model = nn.DataParallel(model, device_ids=[0,1])

其实，将并行后的网络打印出来就会发现需要加上“module”，千万注意是module，而不是model。这样就可以调用并行网络中定义的网络层。

2、在服务器上运行matplotlib库时，因为没有图形界面而报错，具体报错情况没有截图，但是如果使用远程服务器训练，而且没有图形界面。在使用matplotlib绘图时，建议添加以下代码：

import matplotlib.pyplot as plt
plt.switch_backend('agg')

而且训练时不要使用plt.show()，只是保存图片就行，之后在下载就可以。