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
展开阅读全文

没有更多推荐了,返回首页

©️2019 CSDN 皮肤主题: 黑客帝国 设计师: 上身试试
应支付0元
点击重新获取
扫码支付

支付成功即可阅读