PyTorch中常用的工具（5）使用GPU加速：CUDA

前言

在训练神经网络的过程中需要用到很多的工具，最重要的是数据处理、可视化和GPU加速。本章主要介绍PyTorch在这些方面常用的工具模块，合理使用这些工具可以极大地提高编程效率。

由于内容较多，本文分成了五篇文章（1）数据处理（2）预训练模型（3）TensorBoard（4）Visdom（5）CUDA与小结。

整体结构如下：

• 1 数据处理
  • 1.1 Dataset
  • 1.2 DataLoader
• 2 预训练模型
• 3 可视化工具
• 3.1 TensorBoard
• 3.2 Visdom
• 4 使用GPU加速：CUDA
• 5 小结

全文链接：

1. PyTorch中常用的工具（1）数据处理
2. PyTorch常用工具（2）预训练模型
3. PyTorch中常用的工具（3）TensorBoard
4. PyTorch中常用的工具（4）Visdom
5. PyTorch中常用的工具（5）使用GPU加速：CUDA

4 使用GPU加速：CUDA

这部分内容在前面介绍Tensor、nn.Module时已经有所涉及，这里做一个总结，并深入介绍它的相关应用。

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

• Tensor。
• nn.Module（包括常用的layer、损失函数以及容器Sequential等）。

这些数据结构都带有一个.cuda方法，调用该方法可以将它们转为对应的GPU对象。注意，tensor.cuda会返回一个新对象，这个新对象的数据已经转移到GPU，之前的Tensor还在原来的设备上（CPU）。module.cuda会将所有的数据都迁移至GPU，并返回自己。所以module = module.cuda()和module.cuda()效果一致。

除了.cuda方法，它们还支持.to(device)方法，通过该方法可以灵活地转换它们的设备类型，同时这种方法也更加适合编写设备兼容的代码，这部分内容将在后文详细介绍。

nn.Module在GPU与CPU之间的转换，本质上还是利用了Tensor在GPU和CPU之间的转换。nn.Module的.cuda方法是将nn.Module下的所有参数（包括子module的参数）都转移至GPU，而参数本质上也是Tensor。

下面对.cuda方法举例说明，这部分代码需要读者具有两块GPU设备。

注意：为什么将数据转移至GPU的方法叫做.cuda而不是.gpu，就像将数据转移至CPU调用的方法是.cpu呢？这是因为GPU的编程接口采用CUDA，而目前不是所有的GPU都支持CUDA，只有部分NVIDIA的GPU才支持。PyTorch1.8目前已经支持AMD GPU，并提供了基于ROCm平台的GPU加速，感兴趣的读者可以自行查询相关文档。

In: tensor = t.Tensor(3, 4)
    # 返回一个新的Tensor，保存在第1块GPU上，原来的Tensor并没有改变
    tensor.cuda(0)
    tensor.is_cuda # False

Out: False

In: # 不指定所使用的GPU设备，默认使用第1块GPU
    tensor = tensor.cuda()
    tensor.is_cuda # True

Out: True

In: module = nn.Linear(3, 4)
    module.cuda(device = 1)
    module.weight.is_cuda # True

Out: True

In: # 使用.to方法，将Tensor转移至第1块GPU上
    tensor = t.Tensor(3, 4).to('cuda:0')
    tensor.is_cuda

Out: True

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

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

在最后一段代码中，两个Parameter所占用的内存空间都非常大，大约是8GB。如果将这两个Parameter同时放在一块显存较小的GPU上，那么显存将几乎被占满，无法再进行任何其他运算。此时可以通过.to(device_i)将不同的计算划分到不同的GPU中。

下面是在使用GPU时的一些建议。

• GPU运算很快，但对于很小的运算量来说，它的优势无法被体现。因此，对于一些简单的操作可以直接利用CPU完成。
• 数据在CPU和GPU之间的传递比较耗时，应当尽量避免。
• 在进行低精度的计算时，可以考虑使用HalfTensor，它相比于FloatTensor可以节省一半的显存，但是需要注意数值溢出的情况。

注意：大部分的损失函数都属于nn.Module，在使用GPU时，用户经常会忘记使用它的.cuda方法，这在大多数情况下不会报错，因为损失函数本身没有可学习参数（learnable parameters），但在某些情况下会出现问题。为了保险起见，同时也为了代码更加规范，用户应记得调用criterion.cuda，下面举例说明：

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

    # 下面这行会报错，因weight未被转移至GPU
    # loss = criterion(input, target)

    # 下面的代码则不会报错
    criterion.cuda()
    loss = criterion(input, target)

    criterion._buffers

Out: OrderedDict([('weight', tensor([1., 3.], device='cuda:0'))])

除了调用对象的.cuda方法，还可以使用torch.cuda.device指定默认使用哪一块GPU，或使用torch.set_default_tensor_type让程序默认使用GPU，不需要手动调用.cuda方法：

In: # 如果未指定使用哪块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()
        assert a.get_device() == b.get_device() == 1

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

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

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

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

Out: True

如果服务器具有多个GPU，那么tensor.cuda()方法会将Tensor保存到第一块GPU上，等价于tensor.cuda(0)。如果想要使用第二块GPU，那么需要手动指定tensor.cuda(1)，这需要修改大量代码，较为烦琐。这里有以下两种替代方法。

• 先调用torch.cuda.set_device(1)指定使用第二块GPU，后续的.cuda()都无需更改，切换GPU只需修改这一行代码。

• 设置环境变量CUDA_VISIBLE_DEVICES，例如export CUDA_VISIBLE_DEVICE=1（下标从0开始，1代表第二块物理GPU），代表着只使用第2块物理GPU，但在程序中这块GPU会被看成是第1块逻辑GPU，此时调用tensor.cuda()会将Tensor转移至第二块物理GPU。CUDA_VISIBLE_DEVICES还可以指定多个GPU，例如export CUDA_VISIBLE_DEVICES=0,2,3，第1、3、4块物理GPU会被映射为第1、2、3块逻辑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"。如果使用IPython或者Jupyter notebook，那么还可以使用%env CUDA_VISIBLE_DEVICES=1,2设置环境变量。

基于PyTorch本身的机制，用户可能需要编写设备兼容（device-agnostic）的代码，以适应不同的计算环境。在第3章中已经介绍到，可以通过Tensor的device属性指定它加载的设备，同时利用to方法可以很方便地将不同变量加载到不同的设备上。然而，如果要保证同样的代码在不同配置的机器上均能运行，那么编写设备兼容的代码是至关重要的，本节将详细介绍如何编写设备兼容的代码。

首先介绍一下如何指定Tensor加载的设备，这一操作往往通过torch.device()实现，其中device类型包含cpu与cuda，下面举例说明：

In: # 指定设备,使用CPU
    t.device('cpu')
    # 另外一种写法：t.device('cpu',0)

Out: device(type='cpu')

In: # 指定设备，使用第1块GPU
    t.device('cuda:0')
    # 另外一种写法：t.device('cuda',0)

Out: device(type='cuda', index=0)

In: # 更加推荐的做法（同时也是设备兼容的）：如果用户具有GPU设备，那么使用GPU，否则使用CPU
    device = t.device("cuda" if t.cuda.is_available() else "cpu")
    print(device)

Out: cuda

In: # 在确定了设备之后，可以将数据与模型利用to方法加载到指定的设备上。
    x = t.empty((2,3)).to(device)
    x.device

Out: device(type='cuda', index=0)

对于最常见的数据结构Tensor，它封装好的大部分操作也支持指定加载的设备。当拥有加载在一个设备上的Tensor时，通过torch.Tensor.new_*以及torch.*_like操作可以创建与该Tensor相同类型、相同设备的Tensor，举例说明如下：

In: x_cpu = t.empty(2, device='cpu')
    print(x_cpu, x_cpu.is_cuda)
    x_gpu = t.empty(2, device=device)
    print(x_gpu, x_gpu.is_cuda)

Out: tensor([-3.6448e+08,  4.5873e-41]) False
     tensor([0., 0.], device='cuda:0') True

In: # 使用new_*操作会保留原Tensor的设备属性
    y_cpu = x_cpu.new_full((3,4), 3.1415)
    print(y_cpu, y_cpu.is_cuda)
    y_gpu = x_gpu.new_zeros(3,4)
    print(y_gpu, y_gpu.is_cuda)

Out: tensor([[3.1415, 3.1415, 3.1415, 3.1415],
             [3.1415, 3.1415, 3.1415, 3.1415],
             [3.1415, 3.1415, 3.1415, 3.1415]]) False
     tensor([[0., 0., 0., 0.],
             [0., 0., 0., 0.],
             [0., 0., 0., 0.]], device='cuda:0') True

In: # 使用ones_like或zeros_like可以创建与原Tensor大小类别均相同的新Tensor
    z_cpu = t.ones_like(x_cpu)
    print(z_cpu, z_cpu.is_cuda)
    z_gpu = t.zeros_like(x_gpu)
    print(z_gpu, z_gpu.is_cuda)

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

在一些实际应用场景下，代码的可移植性是十分重要的，读者可根据上述内容继续深入学习，在不同场景中灵活运用PyTorch的不同特性编写代码，以适应不同环境的工程需要。

本节主要介绍了如何使用GPU对计算进行加速，同时介绍了如何编写设备兼容的PyTorch代码。在实际应用场景中，仅仅使用CPU或一块GPU是很难满足网络的训练需求的，因此能否使用多块GPU来加速训练呢？

答案是肯定的。自PyTorch 0.2版本后，PyTorch新增了分布式GPU支持。分布式是指有多个GPU在多台服务器上，并行一般指一台服务器上的多个GPU。分布式涉及到了服务器之间的通信，因此比较复杂。幸运的是，PyTorch封装了相应的接口，可以用简单的几行代码实现分布式训练。在训练数据集较大或者网络模型较为复杂时，合理地利用分布式与并行可以加快网络的训练。关于分布式与并行的更多内容将在本书第7章进行详细的介绍。

5 小结

本章介绍了一些工具模块，这些工具有的已经封装在PyTorch之中，有的是独立于PyTorch的第三方模块。这些模块主要涉及数据加载、可视化与GPU加速的相关内容，合理使用这些模块能够极大地提升编程效率。
3,4)
print(y_gpu, y_gpu.is_cuda)

Out: tensor([[3.1415, 3.1415, 3.1415, 3.1415],
[3.1415, 3.1415, 3.1415, 3.1415],
[3.1415, 3.1415, 3.1415, 3.1415]]) False
tensor([[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.]], device=‘cuda:0’) True

```python
In: # 使用ones_like或zeros_like可以创建与原Tensor大小类别均相同的新Tensor
    z_cpu = t.ones_like(x_cpu)
    print(z_cpu, z_cpu.is_cuda)
    z_gpu = t.zeros_like(x_gpu)
    print(z_gpu, z_gpu.is_cuda)

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

在一些实际应用场景下，代码的可移植性是十分重要的，读者可根据上述内容继续深入学习，在不同场景中灵活运用PyTorch的不同特性编写代码，以适应不同环境的工程需要。

本节主要介绍了如何使用GPU对计算进行加速，同时介绍了如何编写设备兼容的PyTorch代码。在实际应用场景中，仅仅使用CPU或一块GPU是很难满足网络的训练需求的，因此能否使用多块GPU来加速训练呢？

答案是肯定的。自PyTorch 0.2版本后，PyTorch新增了分布式GPU支持。分布式是指有多个GPU在多台服务器上，并行一般指一台服务器上的多个GPU。分布式涉及到了服务器之间的通信，因此比较复杂。幸运的是，PyTorch封装了相应的接口，可以用简单的几行代码实现分布式训练。在训练数据集较大或者网络模型较为复杂时，合理地利用分布式与并行可以加快网络的训练。关于分布式与并行的更多内容将在本书第7章进行详细的介绍。