0704-使用GPU加速_cuda

pytorch完整教程目录:https://www.cnblogs.com/nickchen121/p/14662511.html

一、CPU 和 GPU 数据相互转换

在 torch 中以下数据结构分为 CPU 和 GPU 两个版本:

  • Tensor
  • Variable(包括 Parameter)
  • nn.Module(包括常用的 layer、loss function,以及容器 Sequential 等)

它们都带有一个 .cuda 方法,通过这个方法可以把它们转换对应的 GPU 对象。

但是在把 cpu 上的数据转化成 gpu 上的数据时,需要注意以下两点:

  • tensor.cudavariable.cuda 都会返回一个新对象,这个新对象存放在 GPU 中,而之前的数据则依然还会在 CPU 上。
  • module.cuda 会将所有的数据都迁移到 GPU,并且返回自己,也就是说 module=module.cuda()module.cuda() 的效果是一样的

其实 variable 和 nn.Module 在 cpu 和 gpu 之间的转换,本质上还是利用了 tensor 在 cpu 和 gpu 之间的转换。比如 variable.cuda 实际上是把 variable.data 转移到指定的 gpu 上。而 nn.Module 的 cuda 方法是把 nn.Module 下的所有 parameter(包括子 module 的 parameter)都转移到 gpu 上,而 Parameter 的本质其实又是 variable。

下面举例说明,但是需要有两块 gpu 设备。

注:为什么把数据转移到 gpu 的方法叫做 .cuda 而不是 .gpu 呢?这是因为 gpu 的编程接口采用 cuda,而目前并不是所有的 gpu 都支持 cuda,只有部分 NVIDIA 的 gpu 才支持。torch 未来可能还会支持 AMD 的 gpu,而 AMD GPU 的编程接口采用 OpenCL,因此 torch 还预留着 .cl 方法,用于以后支持 AMD 等的 GPU。

import torch as t

# tensor 测试
tensor = t.Tensor(3, 4)
tensor.cuda(0)  # 返回一个新的 tensor,保存在第 1 块 GPU 上,但原来的 tensor 并没有改变
tensor.is_cuda  # False  # 原来的 tensor 依然再 cpu 上

tensor = tensor.cuda()  # 不指定所使用的 GPU 设备,将默认使用第 1 块 GPU
tensor.is_cuda  # False

# variable 测试
variable = t.autograd.Variable(tensor)
variable.cuda()
variable.is_cuda()  # False  # 原来的 variable 依然再 cpu 上

# nn.module 测试
module = nn.Linear(3, 4)
module.cuda(device_id=1)
module.weight.is_cuda  # True

class VeryBigModule(nn.Module):
    def __init__(self):
        super(VeryBigModule, self).__init__()
        self.GiantParameter1 = t.nn.Parameter(t.randn(100000, 20000)).cuda(0)
        self.GiantParameter2 = t.nn.Parameter(t.randn(20000, 100000)).cuda(1)

    def forward(self, x):
        x = self.GiantParameter1.mm(x.cuda(0))
        x = self.GiantParameter2.mm(x.cuda(1))
        return x

在 VeryBigModule 类中,两个 Parameter 所占用的内存非常大,大概是 8GB,如果两者放在一块 GPU 上,可能会把显存占满,因此把这两个 Parameter 放在两块 GPU 上。

二、使用 GPU 的注意事项

关于使用 GPU 有一些小小的建议:

  • gpu 运算很快,但是运算量小时,不能体现出它的优势,因此一些简单的操作可以使用 cpu 完成
  • 数据在 cpu 和 gpu 之间的传递会比较耗时,应当尽量避免
  • 在进行低精度的计算时,可以考虑使用 HalfTensor 时,相比较 FloatTensor 能节省一半的显存,但需要注意数值溢出的情况

注:大部分的损失函数也都属于 nn.Module,但在使用 gpu 时,很多时候我们都忘记使用它的 .cuda 方法,在大多数情况下不会保存,因为损失函数没有可学习的参数。但在某些情况下会出错,为了保险起见也为了代码更规范,也应该记得调用 criterion.cuda,下面举例说明:

# 交叉熵损失函数,带权重
criterion = t.nn.CrossEntropyLoss(weight=t.Tensor([1, 3]))
inp = t.autograd.Variable(t.randn(4, 2)).cuda()
target = t.autograd.Variable(t.Tensor([1, 0, 0, 1])).long().cuda()

# 下面这行会报错,因为 weight 没有被转移到 GPU 上
# loss = criterion(inp, target)

# 这行则不会报错
criterion.cuda()
loss = criterion(inp, target)

criterion._buffers

三、设置默认 GPU

除了调用 .cuda 方法外,还可以使用 torch.cuda.device 指定默认使用哪一块 GPU,或使用 torch.set_default_tensor_type 使程序默认使用 GPU,不需要手动调用 cuda

# 如果没有指定使用哪块 GPU,默认使用 GPU 0
x = t.cuda.FloatTensor(2, 3)
# x.get_device() == 0
y = t.FloatTensor(2, 3).cuda()
# y.get_device() == 0

# 指定默认使用 GPU 1
with t.cuda.device(1):
    # 在 GPU 1 上构建 tensor
    a = t.cuda.FloatTensor(2, 3)

    # 把 tensor 转移到 GPU 1
    b = t.FloatTensor(2, 3).cuda()
    print(a.get_device() == b.get_device() == 1)

    c = a + b
    print(c.get_device() == 1)

    z = x + y
    print(z.get_device() == 0)

    # 手动指定使用 GPU 0
    d = t.randn(2, 3).cuda(0)
    print(d.get_device() == 2)

    # t.set_default_tensor_type('torch.cuda.FloatTensor')  # 指定默认 tensor 的类型为 GPU 上的 FloatTensor
    a = t.ones(2, 3)
    a.is_cuda()

四、GPU 之间的切换

如果服务器有多个 gpu,tensor.cuda() 方法将会把 tensor 保存到第一快 gpu 上,这等价于 tensor.cuda(0),这个时候如果想使用第二块 gpu,需要手动指定 tensor.cuda(1),但是这需要修改大量代码,因此很繁琐。这里有两种代替的方法:

  1. 第一种方法是先调用 t.cuda.set_device(1) 指定使用第二块 gpu,后序的 .cuda() 都不需要改变。
  2. 另外一种方法是设置环境变量 CUDA_VISIBLE_DEVICES,例如当 export CUDA_VISIBLE_DEVICES=1 时,只使用 物理上的第二块 GPU,但在程序中这块 cpu 会被看成是第一块逻辑 gpu。当然,CUDA_VISIBLE_DEVICES 还可以指定多个 gpu,如 export CUDA_VISIBLE_DEVICES=0,2,3,那么第一、三、四块物理 GPU 将会被映射成第一、二、三块逻辑 GPU,也就是说 tensor.cuda(1) 将会把 Tensor 转移到第三块物理 GPU 上。

设置 CUDA_VISIBLE_DEVICES 有两种方法:

  • 一种是在命令行中 CUDA_VISIBLE_DEVICES=0,1 python main.py
  • 另一种是在程序中 import os; os.environ["CUDA_VISIBLE_DEVICES"] = "2"

上述一般都是自己使用的情况,在实际工程中,可能还会用到分布式 GPU,由于一般人员使用不到这种方法,这里不做赘述,想详细了解的可以看官方文档——GPU 分布式通信

内容来源于网络如有侵权请私信删除