深度学习方法——pytorch下GPU与CPU调用的切换

一、准备工作

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

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

输出：
NVIDIA GeForce RTX 3070 Laptop GPU
True

---

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

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

PS：当前博主的电脑只含一块显卡所以无法验证这个代码，大家可以自行测试（参考）

二、修改方法

1. CPU$\to$GPU：使用cuda方法

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

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

输出：
device(type=‘cuda’, index=0)

---

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

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

输出：
cpu cuda:0

---

2. CPU$\leftarrow$GPU：使用cpu方法

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

import torch
x = torch.Tensor([1,2]).cuda() # 生成GPU调用的张量x
y = x.cpu()                    # 修改为CPU调用
print(x.device,y.device)

输出：
cuda:0 cpu

3. GPU$\leftrightarrow$CPU：定义device对象

使用该方法可以通过修改torch.device()中的参数实现GPU和CPU双向切换，因此它具有以下两个优点：

• 在不确定设备是否支持cuda时可以使用条件语句进行定义，从而具有兼容性：

import torch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu' ) # 定义device对象
device

输出：
device(type=‘cuda’)

---

• 该方法可以通过修改device对象，实现多个张量的GPU调用与CPU调用统一定义与统一修改：

import torch
G_device = torch.device('cuda')
# 创建张量时统一定义GPU调用
x = torch.tensor([1,2],device=G_device)
y = torch.tensor([3,4],device=G_device)
z = torch.tensor([5,6],device=G_device)
# 或者采用定义后修改为GPU调用
# x = torch.tensor([1,2]).to(G_device)
# y = torch.tensor([3,4]).to(G_device)
# z = torch.tensor([5,6]).to(G_device)
x.device,y.device,z.device

输出：
(device(type=‘cuda’, index=0),
device(type=‘cuda’, index=0),
device(type=‘cuda’, index=0))

---

修改时使用到的Tensor.to()方法也同样是一个返回新张量的操作，需要注意：

C_device = torch.device('cpu')
# 统一修改为CPU调用
x = x.to(C_device)
y = y.to(C_device)
z = z.to(C_device)
x.device,y.device,z.device

输出：
(device(type=‘cpu’), device(type=‘cpu’), device(type=‘cpu’))

三、补充说明

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