pytorch修改Dataparallel的主gpu

背景

pytorch使用Dataparallel时会出现负载不均衡的现象，一般来说0号gpu占用显存最多。但是，有时候0号gpu显存不是特别多了，但是其他gpu显存很多，更适合用来作为主gpu。所以我们得想办法指定其他gpu为主gpu。

方法

1、修改之前，大家使用多gpu的代码指定device是这样指定的：

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

现在需要修改为：

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

2、修改之前，使用Dataparallel的代码如下：

model = nn.DataParallel(model)

修改之后，加上了device_ids参数的值，并注意，第一个id必须与device中指定的id相同。

model = nn.DataParallel(model, device_ids=[1, 0, 2, 3])

PS: 如果第一个id与device中指定的id不同，例如不指定 device_ids，或指定 device_ids=[0,1,2,3]，都会报如下错误：

RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cuda:1