pytorch中涉及cuda的使用介绍

1. 设置可见GPU的个数

如果你有很多个GPU，而只让程序使用其中的几个GPU，可以使用os方式进行设置。

import os
GPU = '0,1,2' #设置GPU 0 1 2 可见
os.environ['CUDA_VISIBLE_DEVICES'] =GPU

2. GPU的常用操作

2.1 常用操作列表

代码	解释
torch.cuda.is_available()	GPU是否可用(True/False)
torch.version.cuda	查看GPU版本(如：10.2)
torch.cuda.get_device_capability()	查看GPU算力，比如(7,5)
torch.cuda.current_device()	返回当前设备索引(默认为0)
torch.cuda.set_device(1)	设置默认使用的GPU编号
torch.cuda.get_device_name(0)	查看GPU全名
torch.cuda.device_count()	查看可用的GPU

更多关于pytorch如何使用cuda的方法可以参考官网。

2.2 一些操作说明

（1）torch.cuda.is_available()

torch.cuda.is_available()通常的使用方式如下：

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

也可以像下列方式一样将CPU的数据挪到GPU上。

# 也可以是 device = torch.device('cuda:0')
device = torch.device('cuda')
a = torch.tensor([1,2,3])
b = a.to(device )
print(a)
print(b)

输出：

tensor([1, 2, 3])
tensor([1, 2, 3], device='cuda:0')

（2）设置GPU默认使用编号

一般情况下默认使用的GPU是0号，如果你想要改变默认的GPU，可以使用如下方式。

# 里面输入int类型的数字
torch.cuda.set_device(1)

但是我的GPU只有一个卡，所以强制修改默认GPU编号为1，会有如下的报错：

RuntimeError: CUDA error: invalid device ordinal
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

3. 在GPU上使用数据

3.1 直接创建法

可以直接在cuda中创建数据，如下方法所示。

a = torch.ones(3,4,device="cuda")
b = torch.cuda.FloatTensor(3, 4)

3.2 CPU转移至GPU

（1）方法一：tensor.to()

import torch
a = torch.ones(3,4)
b = a.to("cuda")
print(a)
print(b)

输出：

tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]], device='cuda:0')

.to()不仅可以转移device，还可以修改数据类型，比如：a.to(torch.double)。

（2）方法二：tensor.cuda()

a = torch.tensor([1., 2.]).cuda()

（3）方法三：tensor.type()

dtype = torch.cuda.FloatTensor
x = torch.rand(2,2).type(dtype)

（4）方法四：torch.from_numpy(np_labels).cuda()

import torch
import numpy as np

b = np.ones(2)
a = torch.from_numpy(b).cuda()
print(a)

输出：

tensor([1., 1.], device='cuda:0', dtype=torch.float64)

4. 一般哪些对象可以使用GPU

（1）网络模型 (module)

device = torch.device("cuda")
 
module = module.to(device)

（2）数据 (输入imgs、标注target/label)

device = torch.device("cuda")
 
# training
for data in data_loader:
    imgs, targets = data
    imgs = imgs.to(device)
    targets = target.to(device)
    #....
 
# testing/validation
# 同上

（3）损失函数 (loss_function)

device = torch.device("cuda")
 
loss_function = loss_function.to(device)