PyTorch GPU 加速

Pytorch GPU 加速

1. GPU 加速

在 cmd 命令框中输入如下代码，可参看 GPU 的具体信息

nvidia-smi

安装 CUDA，官方地址

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

在 Google Colab 中，需要设置 Runtime -> Change runtime type，在弹出的对话框中，设置 Hardware accelerator 为 GPU。

1.1 常用的 device 相关方法

PyTorch 中 cuda 相关操作的官方教程：网址

torch.cuda.is_available()：返回 bool 类型；GPU 是否可用

torch.device(device): 使用 GPU 或者 CPU 设备

• 参数：device

  • str 类型，{'cpu', 'cuda', 'cuda:0', 'cuda:1', ...}，对于只有一个 GPU 的电脑指定 'cuda' 和 'cuda:0' 等价

  • int 类型， 0, 1, 2, ... 等价于 'cuda:0', 'cuda:1', 'cuda:2'...

device.type 属性：返回 device 的类型

实例 1： 常用设置 device 的代码：如果 GPU 可用，则使用 GPU；如果不可用，使用 CPU

# 如果GPU可用，则使用GPU；如果不可用，使用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

实例 2： 查看 device 相关信息

import torch
use_gpu = torch.cuda.is_available()  # GPU 是否可用
print(torch.cuda.is_available())

print(torch.device('cpu'))      # 获取 CPU 设备
print(torch.device(0))          # 方法 1：获取 GPU 设备
print(torch.device('cuda'))     # 方法 2：获取 GPU 设备
print(torch.device('cuda:0'))   # 方法 3：获取 GPU 设备

gpu_device = torch.device('cuda')    
print(gpu_device, gpu_device.type)  # 参考 GPU 设备类型

if use_gpu:
    print(torch.cuda.device_count())     # GPU 数量
    print(torch.cuda.current_device())   # 当前 GPU 编号
    print(torch.cuda.device(0))          # 获取设备号为 0 的 GPU
    print(torch.cuda.get_device_name(0)) # 获取设备号为 0 的 GPU 名称

1.2 Tensor 数据与 device 相关的属性和方法

Tensor 数据与 device 相关的属性和方法

• torch.tensor(device)：初始化一个 Tensor 数据时，

  参数 device ：str 类型或 int 类型，与上文相同；设置 tensor 数据存放设备（GPU 或 CPU），参数值与上文相同

• Tensor.to(device) 方法：将 tensor 数据复制到指定的设备上

  • model.to(device)方法：类似

• Tensor.cuda(device=None) 方法：将 tensor 数据复制到 GPU 上

  • model.cuda(device) 方法：类似

• Tensor.cpu() 方法：将 tensor 数据复制到 CPU 上

  • model.cpu() 方法：类似

• Tensor.device 属性：可以查看数据所在的设备

实例 1： 设置 tensor 数据的存放设备

import torch
# 定义设备
gpu_device = torch.device("cuda")

# 方法 1：推荐使用，直接在 GPU 上创建
tensor_X1 = torch.tensor([1, 2], device = 0)
tensor_X2 = torch.tensor([1, 2], device = gpu_device)
tensor_X3 = torch.tensor([1, 2], device = "cuda:0")
print(tensor_X1.device, tensor_X2.device, tensor_X3.device)    
# Output:  cuda:0  cuda:0  cuda:0 

# 方法 2：先将 tensor 创建在CPU上，再复制到GPU上
tensor_X1 = torch.tensor([1, 2])       # 在 CPU 上创建 tensor
tensor_X2 = tensor_X1.to(device = gpu_device)
print(tensor_X1.device, tensor_X2.device)
# Output:  cpu  cuda:0

tensor_X3 = tensor_X1.cuda()           # 默认的 device 
tensor_X4 = tensor_X1.cuda(0)          # 其他的参数方式，可以得到相同的结果
tensor_X5 = tensor_X1.cuda(gpu_device) 
tensor_X6 = tensor_X1.cuda("cuda:0")
print(tensor_X3.device, tensor_X4.device, tensor_X5.device, tensor_X6.device)
# Output:  cuda:0  cuda:0  cuda:0  cuda:0

实例 2：将 GPU 上的 tensor 复制到 CPU 上

cpu_device = torch.device("cpu")
# 在 GPU 设备上创建 tensor
tensor_X1 = torch.tensor([1, 2], device=torch.device("cuda")) 
# 将 GPU 上的 tensor 数据复制到 CPU 上
tensor_X2 = tensor_X1.to(cpu_device)
tensor_X3 = tensor_X1.to("cpu")
tensor_X4 = tensor_X1.cpu()
print(tensor_X1.device, tensor_X2.device, tensor_X3.device, tensor_X4.device)    

1.3 使用 GPU 训练

使用 GPU 训练有 2 个要点：

• 将数据放在 GPU 上，有三种方式

  • 在定义 tensor 数据时，通过 device 参数指定
  • .cuda(device) 方法
  • .to(device) 方法

• 将模型参数放在 GPU 上

  • .cuda(device) 方法
  • .to(device) 方法

当训练完成时，如果需要调用数据（模型参数），需要将数据从 GPU 取出到 CPU 上，有两种方法

• .cpu() 方法
• .to(device) 方法

实例：

# 将数据放置于 GPU 上
if use_gpu:
    data = torch.tensor(data, device=gpu_device)  # 方法 1
    data = data.cuda()                            # 方法 2
    data = data.to(device=gpu_device)             # 方法 3

# 将模型参数放置于 GPU 上
if use_gpu:
    # 方式 1
    model = model.cuda()
    loss_fn = loss_fn.cuda()
    # 方式 2
    model = model.to(device=gpu_device)
    loss_fn = loss_fn.to(device=gpu_device)

# 将数据从 GPU 上取出
if not use_gpu:
    # 方式 1
    loss = loss.cpu()
    model.parameters = model.parameters().cpu()
    # 方式 2
    model = model.to(device=cpu_device)
    loss_fn = loss_fn.to(device=cpu_device)

1.4 使用 Dataset 类和 DataLoader 类加载数据

可以在重构 Dataset 类时，在 __getitem__() 方法的最后，将返回的数据加载到 GPU 设备上。

也在 DataLoader 类中可以使用如下代码，将数据加载在 GPU 上。

from torch.utils.data import DataLoader
from torch.utils.data.dataloader import default_collate

dataloader = DataLoader(mydataset,
    collate_fn=lambda x: tuple(x_.to(gpu_device) for x_ in default_collate(x)))  # 将加载的数据置于 GPU 上

2. Colab TPU 加速

参考资料

文中代码：Colab, Github

Pytorch Tutorial, CUDA SEMANTICS, website