Question

相比于对 GPU 完全透明的 Theano 而言，在 PyTorch 中使用 GPU 较为复杂，但同时这也意味着对 GPU 资源更加灵活高效的控制。这部分内容在前面介绍 Tensor、Module 时大都提到过，这里将做一个总结，并深入介绍相关应用。

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

• Tensor
• Variable（包括 Parameter）
• nn.Module（包括常用的 layer、loss function，以及容器 Sequential 等）

它们都带有一个 .cuda 方法，调用此方法即可将其转为对应的 GPU 对象。注意， tensor.cuda 和 variable.cuda 都会返回一个新对象，这个新对象的数据已转移至 GPU，而之前的 tensor/variable 的数据还在原来的设备上（CPU）。而 module.cuda 则会将所有的数据都迁移至 GPU，并返回自己。所以 module = module.cuda() 和 module.cuda() 所起的作用一致。

Variable 和 nn.Module 在 GPU 与 CPU 之间的转换，本质上还是利用了 Tensor 在 GPU 和 CPU 之间的转换。 Variable.cuda 操作实际上是将 variable.data 转移至指定的 GPU。而 nn.Module 的 cuda 方法是将 nn.Module 下的所有 parameter（包括子 module 的 parameter）都转移至 GPU，而 Parameter 本质上也是 Variable。

下面将举例说明，这部分代码需要你具有两块 GPU 设备。

P.S. 为什么将数据转移至 GPU 的方法叫做 .cuda 而不是 .gpu ，就像将数据转移至 CPU 调用的方法是 .cpu ？这是因为 GPU 的编程接口采用 CUDA，而目前并不是所有的 GPU 都支持 CUDA，只有部分 Nvidia 的 GPU 才支持。PyTorch 未来可能会支持 AMD 的 GPU，而 AMD GPU 的编程接口采用 OpenCL，因此 PyTorch 还预留着 .cl 方法，用于以后支持 AMD 等的 GPU。

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

False

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

True

variable = t.autograd.Variable(tensor)
variable.cuda()
variable.is_cuda # False

True

module.cuda?

Object `module.cuda` not found.

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

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

上面最后一部分中，两个 Parameter 所占用的内存空间都非常大，大概是 8 个 G，如果将这两个都同时放在一块 GPU 上几乎会将显存占满，无法再进行任何其它运算。此时可通过这种方式将不同的计算分布到不同的 GPU 中。

关于使用 GPU 的一些建议：

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

另外这里需要专门提一下，大部分的损失函数也都属于 nn.Moudle ，但在使用 GPU 时，很多时候我们都忘记使用它的 .cuda 方法，这在大多数情况下不会报错，因为损失函数本身没有可学习的参数（learnable parameters）。但在某些情况下会出现问题，为了保险起见同时也为了代码更规范，应记得调用 criterion.cuda 。下面举例说明。

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

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

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

criterion._buffers

OrderedDict([('weight', 
               1
               3
              [torch.cuda.FloatTensor of size 2 (GPU 0)])])

而除了调用对象的 .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)

True
True
True
False

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

True

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

• 一种是先调用 t.cuda.set_device(1) 指定使用第二块 GPU，后续的 .cuda() 都无需更改，切换 GPU 只需修改这一行代码。
• 更推荐的方法是设置环境变量 CUDA_VISIBLE_DEVICES ，例如当 export CUDA_VISIBLE_DEVICE=1 （下标是从 0 开始，1 代表第二块 GPU），只使用第二块物理 GPU，但在程序中这块 GPU 会被看成是第一块逻辑 GPU，因此此时调用 tensor.cuda() 会将 Tensor 转移至第二块物理 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" 。如果使用 IPython 或者 Jupyter notebook，还可以使用 %env CUDA_VISIBLE_DEVICES=1,2 来设置环境变量。

从 PyTorch 0.2 版本中，PyTorch 新增分布式 GPU 支持。注意分布式和并行的区别：分布式是指有多个 GPU 在多台服务器上，而并行一般指的是一台服务器上的多个 GPU。分布式涉及到了服务器之间的通信，因此比较复杂，PyTorch 封装了相应的接口，可以用几句简单的代码实现分布式训练。分布式对普通用户来说比较遥远，因为搭建一个分布式集群的代价十分大，使用也比较复杂。相比之下一机多卡更加现实。对于分布式训练，这里不做太多的介绍，感兴趣的读者可参考文档 ^distributed 。