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 = Net()
model = nn.DataParallel(model)
out = model.module.fc(input)

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

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

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

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

其他遇到的bug,会不定时更新~

你可能感兴趣的:(pytorch学习与使用)