设备管理：CPU、GPU 与 torch.accelerator

前面所有例子的张量都待在 CPU 上。真正的深度学习训练离不开 GPU，这一章就讲怎么让代码用上它。

CPU 和 GPU 的分工

CPU 像个全能厨师：啥菜都会做，但一次只能专心做一道。GPU 像流水线上几百个工人：每人只会简单动作，但一起上，几万道工序同时推进。

深度学习的主力运算是大规模矩阵乘法，天生适合”人海战术”。GPU 比 CPU 快几十倍是家常便饭。所以 PyTorch 里每个张量都有归属地——设备（device）：要么在 CPU，要么在 GPU（NVIDIA 显卡上叫 CUDA）。

先看看机器上有什么

import torch

print(torch.cuda.is_available())      # 有没有可用的 NVIDIA GPU
print(torch.cuda.device_count())      # 有几块
print(torch.cuda.get_device_name(0))  # 型号，比如 NVIDIA GeForce RTX 4090

没有 GPU 也没关系，所有代码在 CPU 上照跑不误。用 Mac 的话，Apple 芯片走 mps 设备，用法和 CUDA 几乎一样。这也正是推荐统一写法的原因：

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

张量搬家：to() 一招鲜

张量默认出生在 CPU。想让它去 GPU，用 to()：

x = torch.tensor([1.0, 2.0, 3.0])
print(x.device)   # cpu

x = x.to(device)
print(x.device)   # cuda:0（没 GPU 的话还是 cpu）

x.cuda() 和 x.to("cuda") 是快捷写法，x.cpu() 负责搬回来。也可以在创建时直接指定：

y = torch.randn(3, 4, device=device)

模型也一样，一个 to() 搞定全部参数：

import torch.nn as nn

model = nn.Linear(10, 1).to(device)

模型上了 GPU，喂给它的数据也必须上 GPU。哪一步忘了搬，马上报错。

最经典的报错：设备不一致

a = torch.randn(3).to(device)
b = torch.randn(3)          # 忘了搬，还在 CPU
# c = a + b  # RuntimeError: Expected all tensors to be on the same device

不同设备上的张量不能一起运算。记住两个排查动作：报错后先打印各张量的 .device 看看谁掉队了；训练循环里每个 batch 记得 .to(device)。

GPU 张量要转 NumPy 时也有讲究，得按”摘梯度、回 CPU、转数组”三步走：

arr = x.detach().cpu().numpy()

GPU 上的张量不能直接 .numpy()，带梯度的更不能，顺序别乱。

torch.accelerator：一套代码跑遍所有设备

故事还没完。现在的硬件不止 NVIDIA：苹果的 MPS、英特尔的 XPU、华为昇腾 NPU……每种设备一套专属接口（torch.cuda.*、torch.xpu.*、torch.npu.*），代码写起来很分裂。

PyTorch 2.12 起推出了 torch.accelerator，一个设备无关（device-agnostic）的统一接口。用一套 API 操作”当前加速器”，不管它具体是谁：

import torch

if torch.accelerator.is_available():
    acc = torch.accelerator.current_accelerator()
    print(acc)                                # 比如 cuda 或 xpu
    print(torch.accelerator.device_count())   # 加速器数量
    x = torch.randn(3, 3, device=acc)
    torch.accelerator.synchronize()           # 等加速器上的任务跑完

原来写 torch.cuda.is_available() 的地方，换成 torch.accelerator.is_available() 就通用了。这正是官方把 torch.cuda.* 旧接口逐步向 torch.accelerator 迁移的方向。

显存与多卡，点到为止

GPU 的显存有限，用爆了会报 CUDA out of memory。常用排查手段：减小 batch size、确认验证阶段包了 no_grad()、调用 torch.cuda.empty_cache() 清缓存。另外记录 loss 时一定用 loss.item() 取数值，别让计算图在列表里堆积——这个问题上一章刚讲过。

多卡训练先提一句：老教程里的 nn.DataParallel 写法已不推荐（官方建议用 DistributedDataParallel，简称 DDP），我们会在后面的分布式章节专门讲它。

设备这块就绪之后，下一章开始搭真正的神经网络——先从 nn.Module 这块积木讲起。