3. PyTorch GPU 加速入门

好，我们开始上手。GPU 加速的第一步，就是让 Tensor 从 CPU 搬到 GPU。这一步看似简单，但我见过不少新手在这里翻车——要么忘了迁移，要么用错了方法。

说白了，PyTorch 默认所有 Tensor 都在 CPU 上。你想用 GPU 算，就得手动把它“推”过去。怎么做？两个方法：.cuda() 和 .to()。

3.1 先问问 GPU 在不在

动手之前，先检查机器上有没有可用的 GPU。我个人习惯在训练脚本开头加一段检测：

import torch
print("CUDA available:", torch.cuda.is_available())
print("GPU count:", torch.cuda.device_count())
print("Current device:", torch.cuda.current_device())
print("Device name:", torch.cuda.get_device_name(0))

如果能打出 True 和显卡型号，说明环境准备好了。否则，就算写了 .cuda() 也会报错——我曾经在服务器上忘装驱动，查了一下午才反应过来。

3.2 .cuda() —— 简单粗暴

最直接的方式：

x_cpu = torch.randn(3, 4)
x_gpu = x_cpu.cuda()          # 默认移到 0 号 GPU
x_gpu2 = x_cpu.cuda(1)        # 指定 1 号 GPU

这个方法在早期 PyTorch 里很常见。但你想想看：如果代码里到处写 .cuda()，换设备就得改一大片。我当年第一个项目就是这么干的，后面维护起来真想抽自己。

3.3 .to() —— 更灵活

.to() 可以同时指定设备和数据类型：

device = torch.device("cuda:0")
x_gpu = x_cpu.to(device)

# 或者一步到位
x_gpu = x_cpu.to("cuda:0")
y_gpu = y_cpu.to("cuda:0", dtype=torch.float16)

我个人更推荐用 .to() 方法。因为你可以把 device 定义成一个变量，传给模型和数据，改起来就改一个地方。

两者的区别？我整理了一个对比表：

3.4 设备管理 —— 别再写死 GPU 编号了

项目一上规模，你肯定要面对多 GPU 或者用户机器配置不一致的情况。这时候就得用 torch.device 来做统一管理。

基本用法：

# 在代码开头
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 之后所有 tensor 和模型都往这个 device 上放
model = MyModel().to(device)
data = data.to(device)

这样一来，用户没 GPU 就自动用 CPU，有 GPU 就自动走加速。省心。

多 GPU 时，你可以这样设置主设备：

if torch.cuda.device_count() > 1:
    device = torch.device("cuda:0")
    # 然后用 DataParallel 或者 DistributedDataParallel
    model = nn.DataParallel(model)
    model.to(device)

3.5 避坑指南：我的血泪史

• 迁移后忘了回传： 模型在 GPU 上算完，取结果时用了 .numpy() 直接报错。因为 GPU tensor 不能直接转 numpy，得先 .cpu() 再 .numpy()。
• 不同 GPU 混用： 我曾在两台机器上跑同一套代码，一台单卡一台双卡。因为用了 cuda(0) 写死编号，结果第二台机器主设备不是 0 号，整个 batch 映射乱套。后来改成 torch.device 变量才解决。
• 显存泄漏： 频繁创建 tensor 不及时释放，显存越占越多。用 torch.cuda.empty_cache() 可以清缓存，但最好还是控制好变量生命周期。

3.6 完整流程一览

我画了张图，把整个迁移路径串起来——你看一眼就明白：

从 CPU Tensor 开始，先检测 CUDA 是否可用。能用就通过 .cuda() 或 .to(device) 迁移到 GPU，再统一用设备变量管理多卡，最后计算结果取回 CPU。这条路径我跑了几百次，现在闭着眼也能写对。