Pytorch多GPU训练

参考博客：https://blog.csdn.net/daydayjump/article/details/81158777

使用方法

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训练。

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

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

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

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

model = Net()
model = nn.DataParallel(model)
out = model.module.fc(input)

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