[PyTorch 学习笔记] 7.3 使用 GPU 训练模型

本章代码：

• https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson7/cuda_use.py
• https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson7/multi_gpu.py

这篇文章主要介绍了 GPU 的使用。

在数据运算时，两个数据进行运算，那么它们必须同时存放在同一个设备，要么同时是 CPU，要么同时是 GPU。而且数据和模型都要在同一个设备上。数据和模型可以使用to()方法从一个设备转移到另一个设备。而数据的to()方法还可以转换数据类型。

• 从 CPU 到 GPU

  device = torch.device("cuda")
  tensor = tensor.to(device)
  module.to(device)
  

• 从 GPU 到 CPU

  device = torch.device(cpu)
  tensor = tensor.to("cpu")
  module.to("cpu")
  

  tensor和module的 to()方法的区别是：tensor.to()执行的不是 inplace 操作，因此需要赋值；module.to()执行的是 inplace 操作。

下面的代码是转换数据类型

x = torch.ones((3,3))
x = x.to(torch.float64)

tensor.to() 和 module.to()

首先导入库，获取 GPU 的 device

import torch
import torch.nn as nn
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

下面的代码是执行Tensor的to()方法

x_cpu = torch.ones((3, 3))
print("x_cpu:\ndevice: {} is_cuda: {} id: {}".format(x_cpu.device, x_cpu.is_cuda, id(x_cpu)))

x_gpu = x_cpu.to(device)
print("x_gpu:\ndevice: {} is_cuda: {} id: {}".format(x_gpu.device, x_gpu.is_cuda, id(x_gpu)))

输出如下：

x_cpu:
device: cpu is_cuda: False id: 1415020820304
x_gpu:
device: cpu is_cuda: True id: 2700061800153

可以看到Tensor的to()方法不是 inplace 操作，x_cpu和x_gpu的内存地址不一样。

下面代码执行的是Module的to()方法

net = nn.Sequential(nn.Linear(3, 3))

print("\nid:{} is_cuda: {}".format(id(net), next(net.parameters()).is_cuda))

net.to(device)
print("\nid:{} is_cuda: {}".format(id(net), next(net.parameters()).is_cuda))

输出如下：

id:2325748158192 is_cuda: False
id:1756341802643 is_cuda: True

可以看到Module的to()方法是 inplace 操作，内存地址一样。

torch.cuda常用方法

• torch.cuda.device_count()：返回当前可见可用的 GPU 数量
• torch.cuda.get_device_name()：获取 GPU 名称
• torch.cuda.manual_seed()：为当前 GPU 设置随机种子
• torch.cuda.manual_seed_all()：为所有可见 GPU 设置随机种子
• torch.cuda.set_device()：设置主 GPU 为哪一个物理 GPU，此方法不推荐使用
• os.environ.setdefault(“CUDA_VISIBLE_DEVICES”, “2”, “3”)：设置可见 GPU

在 PyTorch 中，有物理 GPU 可以逻辑 GPU 之分，可以设置它们之间的对应关系。

在上图中，如果执行了`os.environ.setdefault("CUDA_VISIBLE_DEVICES", "2", "3")`，那么可见 GPU 数量只有 2 个。对应关系如下：

逻辑 GPU	物理 GPU
gpu0	gpu2
gpu1</