Pytorch深入浅出（十五）之GPU加速与设备管理

在 PyTorch 中，模型参数与张量默认分配在 CPU 内存中。若希望利用 GPU 的并行计算能力，需要显式地将它们迁移到 GPU 显存。

一、设备迁移

在GPU上运行模型，只需把模型和输入数据转移到GPU上即可 ，其他地方和在CPU上运行无异 。

在 PyTorch 中，将模型或张量转移到指定设备有两种标准写法。

1.cuda()

python 复制代码

if torch.cuda.is_available():
    model = net().cuda()  # 转移模型
    x = x.cuda()          # 转移数据
"""
`.cuda()` 写法的问题在于：
1. 强依赖 CUDA 环境
2. 无法在 CPU / GPU 间无缝切换
3. 不利于部署与调试
"""

2.to.(device)

这是目前工业界的标准写法，它允许你在代码开头统一定义设备，后续代码无需反复修改。

python 复制代码

# 1. 定义设备（Device Agnostic Code）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 2. 迁移模型与数据
model.to(device)
batch_x = batch['img'].to(device)

3.核心避坑：In-place（原地操作）的区别

这是最容易出错的地方，模型和张量在调用迁移方法时的表现截然不同：

对象类型	是否原地操作 (In-place)	正确写法示例
Model (nn.Module)	是	`model.to(device)` （无需重新赋值）
Tensor (张量)	否（返回副本）	`x = x.to(device)` （必须重新赋值）

原理说明：

模型迁移时，PyTorch 会递归地处理所有子层的参数（Parameters），由于这些参数被 Module 持有，可以直接修改。
张量迁移时，数据会从 RAM 被复制到 GPU VRAM 中。原内存中的张量依然存在，返回的是显存中的一个新副本。

4.重要说明

前面讲的两种GPU训练方法：

数据和标签 (data和targets)需要进行 data = data.to(device) 或者 targets = targets.cuda() 赋值，因为它们都是张量，而张量的 .to() 或 .cuda() 默认返回新对象（非原地操作），因此需要赋值
模型和损失函数 作为 nn.Module 的子类，其 .to() 方法被重写为原地操作 ，因此可以直接model.to() model.cuda() loss.to() loss.cuda()而无需赋值

二、多 GPU 环境管理

当你的工作站有多个 GPU 时，PyTorch 默认使用 0 号设备。

1.查询GPU状况

torch.cuda.is_available () -> bool : 检查CUDA是否可用
torch.cuda.device_count () -> int : 返回可用的GPU数量；
torch.cuda.get_device_name (0) -> str : 返回GPU名字，设备索引默认从0开始；
torch.cuda.current_device () : 返回当前设备索引

2.设置环境变量CUDA_VISIBLE_DEVICES

通过屏蔽其他卡，让程序"以为"只有你指定的卡存在。

Jupyter/Python 脚本中（推荐）：

python 复制代码

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2" # 只允许程序看到 1,2 号卡

终端运行项目时：

bash 复制代码

CUDA_VISIBLE_DEVICES=0,1 python train.py

3.逻辑编号重映射（重要规则）

一旦你设置了环境变量，PyTorch 会对可见的 GPU 重新编号。

例如：设置 CUDA_VISIBLE_DEVICES="3"。
此时，程序代码中引用该卡应写 torch.device("cuda:0") 而不是 cuda:3。

三、提取结果：GPU ➡️ CPU ➡️ NumPy

GPU 中的数据不能直接转换为 NumPy 数组，必须先"回家"（回到内存）。

标准三部曲

.cpu() ：将数据从显存移回内存。
.detach() ：将数据从当前计算图中分离（切断梯度追踪）。
.numpy()：转换为 NumPy 格式。

python 复制代码

# 现代规范写法
output = model(batch_input).cpu().detach().numpy()
# 如果数据本来就在 CPU 且不需要梯度，.cpu() 和 .detach() 不会产生额外开销。

四、总结及改进建议

1.强调显存溢出 (OOM) ：在笔记末尾可以提一下，如果报错 RuntimeError: CUDA out of memory，通常是因为 Batch Size 太大或者模型太深，此时需要调小 Batch Size。
2.数据类型一致性 ：迁移到 GPU 后，请确保模型参数和输入数据都是同一种精度（如都是 float32），否则会报类型不匹配错误。