【Pytorch｜CUDA】多卡并行运算时指定主卡

文章目录

• 问题背景
• 多卡并行重新指定主卡
• 参考链接

问题背景

开源代码中使用了多卡并行来计算，我的服务器上有4张显卡，其中0号显卡已经被占用，我只想用 3 号显卡来训练模型，但当我将 id 改为 3 时，提示 0 号显卡上有参数。报错信息如下（当时没有保存，从网上复制一条报错信息，数字可能存在不同）：

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

原来使用多卡并行计算代码时有一张主卡，这张主卡默认为 GPU 0，主卡上会读入参数，然后分发给其他显卡，所以我的目标就是将主卡改为 GPU 3，然后指定在 GPU 3 上进行训练。

多卡并行重新指定主卡

下面代码中在此程序中暂时将GPU3指定为主卡，即id=0，GPU2的id=1，GPU0的id=2，GPU1的id为3。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3,2,0,1"

下面这条命令本来是调用GPU 0,2,3，由于使用了上面的代码给GPU重新编了id号，所以使用的GPU的id是0,2,3，即GPU 3,0,1。

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

参考链接

[1] https://www.jianshu.com/p/536335485aa5