PyTorch中GPU与CPU调用的自由切换

准备工作

在PyTorch中，即使是有GPU的机器，它也不会自动使用GPU，而是需要在程序中显示指定。

在设置GPU调用之前，我们首先需要确认当前设备的显卡是否支持cuda编程，否则所有基于cuda的操作都会报错。我们可以通过torch.cuda.is_available()来进行确认，并通过torch.cuda.get_device_name()查看当前可用的显卡设备名称：

import torch

print(torch.cuda.is_available())     # 验证cuda是否正常安装并能够使用
print(torch.cuda.get_device_name(0)) # 查看使用的设备名称

输出为：

True
NVIDIA GeForce RTX 4090 Laptop GPU

当我们有多块显卡时，可以通过修改系统环境变量中的"CUDA_VISIBLE_DEVICES"，用于指定可以使用的显卡：

import os

os.environ["CUDA_VISIBLE_DEVICES"] = '0'      # 只允许使用0号显卡
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2,3,4'  # 允许使用多块显卡

1、CPU->GPU：使用cuda方法

准备工作完成后，就可以直接在生成张量时使用Tensor.cuda()来将张量切换为GPU调用：

import torch

x = torch.Tensor([1,2]).cuda()
print(x.device)

输出：

cuda:0

注意：cuda()方法并不是一个修改原张量的操作，而是返回一个新的张量，如下所示：

import torch

x = torch.Tensor([1, 2])
y = x.cuda()
print(x.device)
print(y.device)

输出为：

cpu
cuda:0

2、GPU->CPU：使用cpu方法

在张量需要在GPU上执行的操作结束后，我们可以通过Tensor.cpu()将其改回CPU调用的模式，该方法同样是不对原张量进行修改，返回一个新的张量：

import torch

x = torch.Tensor([1, 2]).cuda()
y = x.cpu()
print(x.device)
print(y.device)

输出为：

cuda:0
cpu

3、GPU<->CPU：定义device对象

使用该方法可以通过修改torch.device()中的参数实现GPU和CPU双向切换，因此在不确定设备是否支持cuda时可以使用条件语句进行定义，从而具有兼容性：

import torch

device = torch.device('cuda:0')
# 创建张量时统一定义GPU调用
x = torch.tensor([1, 2], device=device)
y = torch.tensor([3, 4], device=device)
z = torch.tensor([5, 6], device=device)
print(x.device, y.device, z.device)

# 创建张量时统一定义CPU调用
device = torch.device('cpu')
x = torch.tensor([1, 2]).to(device)
y = torch.tensor([3, 4]).to(device)
z = torch.tensor([5, 6]).to(device)
print(x.device, y.device, z.device)

输出为：

cuda:0 cuda:0 cuda:0
cpu cpu cpu

修改device时使用到的Tensor.to()方法也同样是一个返回新张量的操作，如下所示：

import torch

device = torch.device('cuda:0')

x = torch.tensor([1, 2])
y = torch.tensor([3, 4])
z = torch.tensor([5, 6])

x = x.to(device)
y = y.to(device)
z = z.to(device)
print(x.device, y.device, z.device)

输出为：

cuda:0 cuda:0 cuda:0

4、补充说明

上述的Tensor.cuda()、Tensor.cpu()、Tensor.to()等方法在张量（Tensor）中使用时就是返回一个新的张量。但是，如果是由模型（Model）对象使用这些方法，即：model.cuda()、model.cpu()、model.to()，就会直接对原模型产生影响，如下所示：

import torch
import torch.nn as nn

device = torch.device('cuda:0')

model1 = nn.Linear(20, 30)
# 默认存放在GPU
print(next(model1.parameters()).device)

# 直接对原模型产生影响
model2 = model1.cuda()
print(next(model1.parameters()).device)
print(next(model2.parameters()).device)

输出为：

cpu
cuda:0
cuda:0

5、参考 

PyTorch中GPU与CPU调用的自由切换 - 飞燕网