PyTorch学习—21.GPU的使用

引言

  本节学习如何使用GPU进行加速模型运算，介绍Tensor和Module的to函数使用以及它们之间的差异，同时学习多GPU运算的分发并行机制。

一、CPU与GPU

CPU（Central Processing Unit, 中央处理器）：主要包括控制器和运算器
GPU(Graphics Processing Unit, 图形处理器)：处理统一的，无依赖的大规模数据运算

  处理器处理数据的运算，数据必须位于同一个处理器上，要么，同时存在CPU上；要么，同时存在GPU上。那么，在PyTorch中，我们如何将数据在CPU与GPU上进行切换呢？PyTorch中的to函数实现了在CPU与GPU之间的切换。

# CPU->GPU
data.to("cuda")

# GPU->CPU
data.to("cpu")

PyTorch中有两种数据类型：

• Tensor
• Module

  针对这两种数据类型，都有to函数（转换数据类型/设备），下面我们来学习这两种数据类型中的to函数。

• tensor.to(*args, **kwargs)

  # 转换数据类型
  x = torch.ones((3, 3))
  x = x.to(torch.float64)
  
  # 转换数据设备
  x = torch.ones((3, 3))
  x = x.to("cuda")
  

• module.to(*args, **kwargs)

  # 转换模型中数据的类型-所有数据
  linear = nn.Linear(2, 2)
  linear.to(torch.double)
  
  # 转换模型数据设备
  gpu1 = torch.device("cuda")
  linear.to(gpu1)
  

Tensor与 Module的to函数的区别：张量不执行inplace（原地操作），即to函数后会构建新的张量；模型执行inplace（原地操作）。这也就是为什么 Module的to函数不需要“=”重新赋值。下面，我们通过代码来学习这两个方法：

Tensor的to函数:

import torch
import torch.nn as nn

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


# ========================== tensor to cuda
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: 1515330671360
x_gpu:
device: cuda:0 is_cuda: True id: 1515354356800

可以发现：Tensor的to函数非原地操作。
Module的to函数:

import torch
import torch.nn as nn

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

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:2807152911216 is_cuda: False
id:2807152911216 is_cuda: True

可以发现：Module的to函数原地操作。
在GPU上运行，

import torch
import torch.nn as nn


x_cpu = torch.ones((3, 3))
x_gpu = x_cpu.to(device)

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

output = net(x_gpu)
print("output is_cuda: {}".format(output.is_cuda))

output is_cuda: True