PyTorch 笔记（19）— Tensor 用 GPU 加速

2022-09-22 12:55:35

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

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

它们都带有一个.cuda 方法，调用此方法即可将其转为对应的GPU 对象。

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

tensor= t.Tensor(3,4)
# 返回一个新的tensor，但原来的tensor并没有改变
tensor.cuda(0)
tensor.is_cuda  # False

重新赋给自己，tensor 指向cuda 上的数据，不再执行原数据。不指定使用的GPU 设备，将默认使用第 1 块GPU。

tensor= tensor.cuda()
tensor.is_cuda	# True

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

# 将nn.Module模型放到cuda上，其子模型也都自动放到cuda上
from torchimport nn
module= nn.Linear(3,4)
module.cuda(device=0)
module.weight.is_cuda	# True

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

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

classVeryBigModule(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)defforward(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 能节省一半的显存，但需千万注意数值溢出的情况；

而除了调用对象的.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 0with t.cuda.device(0):# 在GPU 0上构建tensor
    a= t.cuda.FloatTensor(2,3)# 将tensor转移至GPU 0
    b= t.FloatTensor(2,3).cuda()print(a.get_device()== b.get_device()==0)# True

    c= a+ bprint(c.get_device()==0)# True

    z= x+ yprint(z.get_device()==0)# True# 手动指定使用GPU 0
    d= t.randn(2,3).cuda(0)print(d.get_device()==2)# False
# 指定默认tensor的类型为GPU上的FloatTensor
t.set_default_tensor_type('torch.cuda.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 来设置环境变量。

从 0.4 版本开始，PyTorch 新增了tensor.to(device) 方法，能够实现设备透明，便于实现CPU/GPU 兼容。

x= torch.randn(1)if torch.cuda.is_available():
    device= torch.device("cuda")            
    y= torch.ones_like(x, device=device)# directly create a tensor on GPU
    x= x.to(device)                         
    z= x+ yprint(z)print(z.to("cpu", torch.double))

输出：

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

PyTorch 支持分布式GPU 。分布式是指有多个GPU 在多台服务器上，而并行一般指的是一台服务器上的多个GPU。分布式涉及到了服务器之间的通信，因此比较复杂，PyTorch 封装了相应的接口，可以用几句简单的代码实现分布式训练。分布式对普通用户来说比较遥远，因为搭建一个分布式集群的代价十分大，使用也比较复杂。相比之下一机多卡更加现实。

  • 作者：wohu1104
  • 原文链接：https://blog.csdn.net/wohu1104/article/details/107239936
    更新时间：2022-09-22 12:55:35