GPU训练实战：用类的__call__方法封装优雅的PyTorch训练器

1. 为什么第37天我决定把"GPU训练"和"类的call方法"放在一起写

先交代一下背景。这是我记录深度学习实战笔记的第37天，前36天里我一直在折腾环境配置、数据管线、模型调参，踩过的坑能装满一个集装箱。今天这个主题看起来有点怪——"GPU训练"和"Python类的__call__方法"有什么关系？一个偏硬件一个偏语法，强行放在一起是不是有点凑篇幅？

还真不是。

如果你写过一段时间的训练脚本，你一定遇到过这种情况：代码越写越长，训练流程从"跑通就行"变成了"要能复现、要能换数据集、要能调超参、要能断点续训"，然后你开始把所有参数攒成一个argparse，把数据加载、模型初始化、训练循环、验证循环全部塞进一个几百行的main()函数里。这个函数改一次崩一次，换一个实验配置就等于重新梳理一遍逻辑。这时候你会发现，训练一个模型的本质，就是不断调用一个有状态的对象——它记住了模型参数，记住了优化器状态，记住了当前epoch，你只需要告诉它"继续"或者"再来一轮"。

而Python里恰好有一个专门为这种场景设计的语法特性，就是__call__。

所以这篇文章我把两件事串起来讲：第一，GPU训练环境怎么正确配置、怎么验证、怎么避开那些让人崩溃的坑；第二，怎么用类的__call__方法把训练流程封装成一个优雅的、可复用的"训练器"对象。这一天的笔记写完之后，我个人是觉得这俩知识点放在一起，比单拎出来任何一个都有价值——因为GPU解决的是"算得快"的问题，而类的设计解决的是"写得顺"的问题，一个完整的深度学习项目，这两件事缺一不可。

2. 先搞定GPU训练环境：从驱动到PyTorch的全链路验证

很多人一上来就问"为什么我的PyTorch用不了GPU"，然后就开始重装CUDA、重装驱动，折腾一整天最后还是不行，其实根子上就是没搞明白GPU训练环境的层级关系。这一节我把完整的链路捋一遍，你照着走一遍，基本不会再出幺蛾子。

2.1 GPU训练环境的三层结构：驱动、CUDA Runtime、深度学习框架

GPU训练的环境从来不是"装一个东西就能跑"，它是三层结构叠加的：

注意一个关键点：PyTorch这种框架，它内部的CUDA运行库是自带的，不依赖你系统里单独装的CUDA Toolkit。也就是说，你装PyTorch GPU版时，它会把对应版本的CUDA runtime一起打包进来。你系统里装不装CUDA Toolkit，其实都不影响PyTorch跑GPU训练——前提是显卡驱动版本足够新。

这就是为什么很多人装了CUDA Toolkit反而把环境搞乱了。你只需要做两件事：

1. 安装一个足够新的NVIDIA显卡驱动；
2. 用PyTorch官方推荐的pip install方式安装GPU版PyTorch。

就这么简单。

2.2 验证GPU是否真正可用的完整步骤

环境装好之后，不要急着跑训练，先做三轮验证。我每次换机器、换环境都要走一遍这三步，每一步都有它的意义。

第一步：验证驱动是否正常。

bash复制nvidia-smi

如果能正常打印出显卡型号、驱动版本、显存占用，说明第一层没问题。这个命令输出的右上角有一个"CUDA Version"字段，注意它指的是当前驱动支持的最高CUDA版本，不是系统里装的CUDA版本。比如驱动显示CUDA Version: 12.4，那你在PyTorch里装CUDA 12.4或更低版本都没问题。

常见故障：如果nvidia-smi报错"无法与NVIDIA驱动通信"或者"NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver"，那说明驱动装坏了。在Windows上还有一个更隐蔽的坑——通过Windows Update自动更新的驱动，有时候会被系统策略锁住，设备管理器里显示"Windows仍在设置此设备的类配置，代码56"，这种情况建议进安全模式用DDU完全卸载驱动，然后重新安装NVIDIA官网的正式版驱动。

第二步：验证PyTorch能否正常调用GPU。

python复制import torch
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.device_count())
print(torch.cuda.get_device_name(0))

这几行输出里，torch.cuda.is_available()返回True才算真正能用。如果返回False，要么是PyTorch装成了CPU版本，要么是驱动版本太旧不满足要求。

有个很常见的翻车现场：用conda或pip装了torch之后，is_available()返回False，但是nvidia-smi完全正常。这时候十有八九是装成了CPU版的PyTorch。我遇到过最夸张的一次，是一个同学用国内镜像源装的torch，镜像同步滞后，拉下来的还是CPU版本的包。验证方法是看torch.__version__后面有没有+cu118、+cu121这类后缀，纯数字版本号基本就是CPU版。

第三步：做一个真实的小矩阵运算验证性能。

python复制import torch
import time

x = torch.randn(10000, 10000)
y = torch.randn(10000, 10000)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x, y = x.to(device), y.to(device)

torch.cuda.synchronize()
start = time.time()
z = x @ y
torch.cuda.synchronize()
print(f"计算耗时: {time.time() - start:.4f}秒")
print(f"结果Shape: {z.shape}")

注意这里我加了torch.cuda.synchronize()，这是GPU编程里一个非常重要的细节。GPU计算是异步的——你调用x @ y，Python这行代码会立即返回，但GPU可能还在后台计算。如果直接用time.time()测耗时，大概率测出来是接近0的假数据。torch.cuda.synchronize()的作用是阻塞当前线程，等GPU上的计算全部完成再往下走，这样才能测到真实耗时。

2.3 多个GPU同时工作时的正确姿势

很多人第一次用多卡训练，都是在服务器上——机器上有8张卡，想同时跑几个不同的实验。这里有一个热词我印象很深，叫"Linux三个GPU同时测试"。

最直接的方式是环境变量控制可见设备：

bash复制# 只让程序看到 GPU 0 和 GPU 1
CUDA_VISIBLE_DEVICES=0,1 python train.py

# 只让程序看到 GPU 2
CUDA_VISIBLE_DEVICES=2 python train.py

在Python代码里也可以用os.environ设置：

python复制import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

这个环境变量的作用是物理屏蔽——设置之后，程序里torch.cuda.device_count()只返回可见设备的数量，cuda:0对应的是你列表里的第一张卡。这样写的好处是，代码里不用硬编码设备编号，换卡的时候只改环境变量就行。

还有一个很实用的小技巧：如果你想每张卡跑一个独立实验，用nohup或者tmux分别启动多个进程，每个进程指定不同的CUDA_VISIBLE_DEVICES，互不干扰。GPU的算力独占没有太大问题，最怕的是显存爆掉——nvidia-smi能看到的"Memory-Usage"如果长期接近100%，说明有人在抢显存。

3. 训练脚本烂摊子：为什么我最后转向了类封装

环境搞定了，下一步就该写训练代码了。我先给你看看我以前写训练脚本的"坏味道"，你对照一下自己有没有中招。

3.1 一个典型的过程式训练代码有多痛苦

早期写训练代码，基本长这样：

python复制import torch
import torch.nn as nn
from torch.utils.data import DataLoader

def train(model, train_loader, val_loader, epochs, lr, device):
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    
    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()