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)方法
- 在定义 tensor 数据时，通过
-
将模型参数放在 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 加速
参考资料
Pytorch Tutorial, CUDA SEMANTICS, website
