在GPU上运行pytorch程序（指定单/多显卡)

目录

1. 使用CUDA_VISIBLE_DEVICES

2. 使用cuda()和torch.cuda.set_device()

3. 使用device

4. 使用torch.nn.DataParallel

---

1. 使用CUDA_VISIBLE_DEVICES

使用CUDA_VISIBLE_DEVICES设置显卡https://blog.csdn.net/qq_43307074/article/details/127659967

2. 使用cuda()和torch.cuda.set_device()

torch.cuda常用指令https://blog.csdn.net/qq_43307074/article/details/127628498?spm=1001.2014.3001.5501方法1和方法2可以同时使用，比如在运行代码时使用：

CUDA_VISIBLE_DEVICES=2,3 python test.py

而在代码内部又指定：

model.cuda(1)

那么model会在GPU3上运行。原理是CUDA_VISIBLE_DEVICES遍历当前可见的设备，并从零开始为可见设备编号。CUDA_VISIBLE_DEVICES使得只有GPU2,3可见，程序就会把这两张显卡编号为GPU0,1,2,3，cuda(1)把model加载到了GPU1上，则实际使用的显卡是GPU3。

如果利用.cuda()或torch.cuda.set_device()把模型加载到多个显卡上，而实际上只使用一张显卡运行程序的话，那么程序会把模型加载到第一显卡上，在运行代码时使用：

CUDA_VISIBLE_DEVICES=2,3 python test.py

而在代码内部又指定：

model.cuda('cuda:1,0')

那么model会在GPU3上运行。

3. 使用device

3.1. 检查所用的device

import torch
x = torch.tensor([[1,2,3],[4,5,6]])
print(x.device)

3.2. 模型/数据指定GPU，具体有以下几种形式：

import torch
cuda = torch.device('cuda:0'/'cuda')
x = torch.tensor([1,2,3],device = cuda)
print(x.device)

import torch
device = torch.device('cuda:0'/'cuda')
x = torch.rand((4,5)).to(device)
print(x.device)

import torch
x = torch.tensor([1,2,3],device = torch.device('cuda:0'/'cuda')
print(x.device)

import torch
device = torch.device('cuda:0'/'cuda')
...
net = net.to(device)
print(net.device)

 经典写法：

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

如果想要在其他GPU上运行，举例如下：

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

cuda = torch.device('cuda:1')
x = torch.tensor([[1,2,3]],device = cuda)
print(x.device)

4. 使用torch.nn.DataParallel

多卡数据并行一般使用torch.nn.DataParallel

torch.nn.DataParallel(model,device_ids)
举例：
torch.nn.DataParallel(model, device_ids=device_ids)
torch.nn.DataParallel(modul, device_ids=[x1,x2,x3,x4,x5,x6,x7])
# 使用的GPU一定是编号连续的
torch.nn.DataParallel(model,device_ids = range(torch.cuda.device_count()) )

其中model是需要运行的模型，device_ids指定部署模型的显卡，数据类型是list/device。

device_ids中的第一个GPU（即device_ids[0]）和model.cuda()或torch.cuda.set_device()中的第一个GPU序号应保持一致，否则会报错。此外如果两者的第一个GPU序号都不是0,比如设置为：

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

那么程序可以在GPU2和GPU3上正常运行。device_ids的默认值是使用可见的GPU，不设置model.cuda()或torch.cuda.set_device()等效于设置了model.cuda(0)

具体举例如下：

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

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

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

所有的数据要先load到主GPU上，然后再分发给每个GPU去train，因此应当将model先放到GPU上，然后在进行并行训练。

参考

在pytorch中指定显卡https://zhuanlan.zhihu.com/p/166161217