在单机多gpu环境下使用pytorch指定gpu

在单机多gpu环境下使用pytorch指定gpu

os.environ['CUDA_VISIBLE_DEVICES'] = '1,2,3' 这条代表表示可见的gpu编号，本机编号从0开始，现在只有1,2,3号卡空闲
在一个4卡机上进行测试，发现设置os.environ无论如何gpu可用设备都有4台
print('available gpus is ', torch.cuda.device_count(), torch.cuda.get_device_name())
此时只需要将os.environ的位置放置import os之前，import torch之后，当前设备就变成了3

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'
import torch

此时我们在进行多gpu训练,注意pytorch中认为你只有3个设备且gpu编号为[0,1,2]对应真实机器的gpu编号为[1,2,3]也就是你在os.environ中设置的顺序。

多gpu代码测试：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'
import torch
use_gpu = [0,2] # 表示我想使用可用设备中的0,2号机器，对应真实gpu编号为1,3
if torch.cuda.device_count() > 1:
    print('available gpus is ', torch.cuda.device_count(), torch.cuda.get_device_name())
    rnn = torch.nn.DataParallel(rnn, use_gpu)
rnn = rnn.cuda()