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()
            optimizer.step()
            total_loss += loss.item()
        
        # 验证
        model.eval()
        val_loss = 0
        correct = 0
        with torch.no_grad():
            for data, target in val_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                val_loss += criterion(output, target).item()
                pred = output.argmax(dim=1)
                correct += pred.eq(target).sum().item()
        
        print(f"Epoch {epoch}: train_loss={total_loss/len(train_loader):.4f}, "
              f"val_loss={val_loss/len(val_loader):.4f}, "
              f"acc={correct/len(val_loader.dataset):.4f}")

这段代码看着没问题，跑个简单分类任务完全够用。但问题在于：一旦你开始做实验，这段代码的每一个环节都要改。

比如你要加一个学习率调度器，你得在train函数里多传一个参数；你要加混合精度训练，得改前向传播和反向传播的部分；你要做梯度累积，得改epoch循环里的逻辑；你要记录各种指标到TensorBoard，得在训练循环里加日志代码；你要定期保存checkpoint，还得加文件保存逻辑。最后这个函数会膨胀到几百行，参数列表长到能写三屏，每个新实验都是从旧代码里复制粘贴再改，改错一个地方就浪费半天。

更麻烦的是模型自身的逻辑和训练逻辑耦合在了一起。你今天用ResNet，明天换ViT，再后天又换成自己改的网络结构，但是训练流程的代码基本上是重复的。你真正需要做的，是让"训练"这件事变成一个可以被反复调用的组件——传入一个模型、一份数据、一堆超参，它就帮你完成训练的完整流程。

3.2 从面向过程到面向对象的思维转变

那怎么解决？核心思路就是把"训练"这件事抽象成一个类。

其实仔细想想，训练过程天然就是有状态的：它需要记住当前的epoch数、当前的优化器状态、当前的best accuracy、当前的模型参数。这些状态如果散落在函数体里，只能靠函数的局部变量和返回值传来传去，非常别扭。但如果你用一个对象去承载这些状态，一切都变得自然了。

这就是类的天然优势。类是"状态 + 行为"的封装。状态是你训练过程中的所有变量，行为是train一个epoch、validate、保存checkpoint这些操作。你把这些东西封装好之后，调用方只需要做一件事：创建一个训练器对象，然后调用它。

而__call__方法的加入，让"调用"这件事变得极其优雅。

4. Python 类 __call__ 方法的本质与使用场景

在进入最终的训练器设计之前，我得先把__call__这个方法彻底讲透。因为这个知识点如果你只是知道"能调用实例"，那概念是记不住的，你得理解它的设计意图。

4.1 __call__ 让实例像函数一样被调用

在Python里，任何对象都可以被"调用"——只要你给它加上一对括号。函数可以调用，类可以调用（创建实例），那实例本身能不能调用？

默认情况下不能。比如你写：

python复制class Person:
    def __init__(self, name):
        self.name = name

p = Person("Tom")
p()  # TypeError: 'Person' object is not callable

会报错，因为Person这个类没有实现__call__方法。但如果加上它：

python复制class Person:
    def __init__(self, name):
        self.name = name
    
    def __call__(self, greeting="Hello"):
        return f"{greeting}, {self.name}"

p = Person("Tom")
print(p())          # Hello, Tom
print(p("Hi"))      # Hi, Tom
print(p.__call__("Hey"))  # 等价写法

这时候p就变成可调用对象了。你调用p()的时候，Python解释器会在背后调用p.__call__()。这就是__call__的工作原理——语法糖而已，但用起来非常自然。

很多人拿这个和普通方法对比，问有什么区别。区别在于：普通方法你要写p.say("Hi")，而__call__让你能直接写p("Hi")。前者是"给这个对象发一个名为say的消息"，后者是"把这个对象当成一个函数用"。从语义上讲，后者更强烈地暗示了"这个东西的主要用途就是被调用"。

4.2 __call__ 和普通方法、静态方法、类方法的对比

为了让你彻底搞清楚__call__在整个类方法体系里的位置，我放一张对比：

一个关键理解：__call__是协议（Protocol），不是普通的命名方法。Python里有一批这样的"双下划线"方法（也叫dunder方法），它们定义了对象的语言级行为。__len__定义len(obj)的行为，__iter__定义for x in obj的行为，__call__定义obj()的行为。你要做的就是实现这些方法，Python会在特定语法场景下自动调用它们。

4.3 为什么要用 __call__ 而不是普通方法

用obj.__call__()和用obj.train()其实都能完成功能，那我为什么推荐用__call__？有这几个理由：

第一，语义清晰。 一个类如果实现了__call__，那它就是在告诉使用者："这个类的对象就是用来被调用的。"你不用去查文档里方法名是train还是run还是fit还是execute——直接obj(...)就行，这是Python语言层面的统一约定。

第二，与函数式编程范式兼容。 Python里很多高级用法——map、filter、装饰器、functools.partial——都要求你传入一个"可调用对象"。类是callable，函数也是callable，它们可以无缝对接。

第三，便于携带状态。 这一点在训练场景里特别重要。一个函数如果要携带中间状态，你得用全局变量、闭包或者额外参数。但一个实现__call__的类对象，它的状态就挂在self上——上一次调用的结果、累计的统计信息、内部的计数器，随时可以访问和修改。

这里有一个很典型的例子，就是functools.partial。它返回的就是一个类似__call__的对象，底层实现就是实现了__call__的类。每次你对一个函数做部分参数绑定，得到的那个"新函数"，本质就是一个可调用对象。

5. 实战：用 __call__ 设计一个优雅的GPU训练器

基础理论讲完了，现在进入正题——把前面所有内容串起来，写一个真正能用的训练器。

5.1 训练器类的整体设计思路

我先说设计目标，你对照着看代码：

1. 初始化时传入所有配置（模型、数据、超参、设备等），不需要每次调用再传一堆参数；
2. 实现__call__方法，调用一次就是跑完整训练流程（或者跑一个epoch，看你的设计）；
3. 内部管理训练状态（当前epoch、best metric、优化器、调度器等）；
4. 支持常见的训练增强逻辑（混合精度、梯度裁剪、checkpoint保存、早停，至少留接口）；
5. 对GPU友好：自动检测设备、显式管理模型和数据在设备间的移动。

我把代码分成几块来讲，先看核心骨架：

python复制import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.cuda.amp import GradScaler, autocast
import time
import json
from pathlib import Path


class GPUTrainer:
    """一个用 __call__ 封装 GPU 训练流程的训练器类"""
    
    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        criterion: nn.Module,
        optimizer: torch.optim.Optimizer,
        device: str = None,
        epochs: int = 10,
        use_amp: bool = False,          # 混合精度训练
        grad_clip: float = None,        # 梯度裁剪阈值
        save_dir: str = "checkpoints",  # 模型保存目录
        scheduler=None,                 # 学习率调度器
        callbacks=None,                 # 回调函数列表
    ):
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.criterion = criterion
        self.optimizer = optimizer
        self.scheduler = scheduler
        
        # 设备自动选择
        if device is None:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.device = torch.device(device)
        
        self.epochs = epochs
        self.use_amp = use_amp
        self.grad_clip = grad_clip
        self.callbacks = callbacks or []
        
        # 训练状态
        self.current_epoch = 0
        self.best_metric = float("-inf")
        self.history = []
        
        # 检查点目录
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
        
        # 混合精度scaler
        self.scaler = GradScaler(enabled=use_amp)
        
        # 把模型放到设备上
        self.model = self.model.to(self.device)

5.2 每个epoch训练与验证的具体实现

然后是fit方法——这里我用一个内部方法实现，不放__call__里。__call__的定位是"启动训练"的入口，具体的单轮训练逻辑放到_train_epoch和_validate里。

python复制    def _train_epoch(self):
        """执行一个epoch的训练，返回平均loss"""
        self.model.train()
        total_loss = 0
        n_batches = 0
        
        for batch_idx, (data, target) in enumerate(self.train_loader):
            data, target = data.to(self.device), target.to(self.device)
            
            self.optimizer.zero_grad()
            
            if self.use_amp:
                # 混合精度正向传播
                with autocast():
                    output = self.model(data)
                    loss = self.criterion(output, target)
                
                # 反向传播
                self.scaler.scale(loss).backward()
                
                # 梯度裁剪（如果有）
                if self.grad_clip is not None:
                    self.scaler.unscale_(self.optimizer)
                    nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip)
                
                self.scaler.step(self.optimizer)
                self.scaler.update()
            else:
                # 普通精度训练
                output = self.model(data)
                loss = self.criterion(output, target)
                loss.backward()
                
                if self.grad_clip is not None:
                    nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip)
                
                self.optimizer.step()
            
            total_loss += loss.item()
            n_batches += 1
            
            # 打印进度
            if (batch_idx + 1) % 50 == 0:
                print(f"  Batch {batch_idx + 1}/{len(self.train_loader)}, "
                      f"Loss: {loss.item():.4f}")
        
        if self.scheduler is not None:
            self.scheduler.step()
        
        return total_loss / n_batches
    
    def _validate(self):
        """执行验证，返回平均loss和准确率"""
        self.model.eval()
        val_loss = 0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for data, target in self.val_loader:
                data, target = data.to(self.device), target.to(self.device)
                
                if self.use_amp:
                    with autocast():
                        output = self.model(data)
                        loss = self.criterion(output, target)
                else:
                    output = self.model(data)
                    loss = self.criterion(output, target)
                
                val_loss += loss.item()
                pred = output.argmax(dim=1)
                correct += pred.eq(target).sum().item()
                total += target.size(0)
        
        return val_loss / len(self.val_loader), correct / total

5.3 核心：__call__ 方法把整个训练流程串起来

重点来了。__call__实现整个训练流程的调度，同时管理checkpoint保存：

python复制    def __call__(self, trained_epochs: int = None):
        """调用对象即开始训练。
        
        Args:
            trained_epochs: 支持只训练指定epoch数，默认使用初始化时的epochs
        """
        target_epochs = trained_epochs or self.epochs
        
        print(f"===== 开始训练 =====")
        print(f"设备: {self.device}")
        print(f"模型参数量: {sum(p.numel() for p in self.model.parameters()) / 1e6:.2f}M")
        print(f"训练集大小: {len(self.train_loader.dataset)}")
        print(f"验证集大小: {len(self.val_loader.dataset)}")
        
        for epoch in range(self.current_epoch, target_epochs):
            self.current_epoch = epoch + 1
            start_time = time.time()
            
            train_loss = self._train_epoch()
            val_loss, val_acc = self._validate()
            elapsed = time.time() - start_time
            
            # 记录历史
            record = {
                "epoch": self.current_epoch,
                "train_loss": train_loss,
                "val_loss": val_loss,
                "val_acc": val_acc,
                "elapsed_sec": elapsed,
            }
            self.history.append(record)
            
            print(f"Epoch {self.current_epoch}/{target_epochs} | "
                  f"train_loss: {train_loss:.4f} | "
                  f"val_loss: {val_loss:.4f} | "
                  f"val_acc: {val_acc:.4f} | "
                  f"time: {elapsed:.1f}s")
            
            # 保存最佳模型
            if val_acc > self.best_metric:
                self.best_metric = val_acc
                self.save_checkpoint("best_model.pth")
            
            # 执行回调
            for callback in self.callbacks:
                callback(self)
        
        print(f"===== 训练完成，最佳验证准确率: {self.best_metric:.4f} =====")
        return self.best_metric

调用方式极其简洁：

python复制trainer = GPUTrainer(
    model=model,
    train_loader=train_loader,
    val_loader=val_loader,
    criterion=nn.CrossEntropyLoss(),
    optimizer=torch.optim.Adam(model.parameters(), lr=1e-3),
    epochs=10,
    use_amp=True,
    device="cuda",
)
trainer()  # 这一句话，就完成了整个训练

如果断点续训或者只跑几轮看看效果，直接：

python复制trainer(3)  # 只训练到第3个epoch，模型和优化器状态还是保留的
trainer()   # 继续跑到初始化时的10个epoch

这个体验是不是跟fit、train这类方法完全不一样？你拿到的是一个"可调用的训练器对象"，它自己记住了所有状态，你只需要像调用函数一样启动它。

5.4 保存和加载Checkpoint的细节

GPU训练里最怕的就是训练到一半机器挂了、显存炸了、或者你发现学习率设置不对想重来。所以checkpoint机制必不可少。这里我给出的保存内容不仅是模型权重，还包括优化器、scheduler、epoch和best metric，这样能真正实现"无缝恢复"：

python复制    def save_checkpoint(self, filename: str):
        """保存完整训练状态，而不仅仅是模型权重"""
        checkpoint = {
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "scheduler_state_dict": self.scheduler.state_dict() if self.scheduler else None,
            "epoch": self.current_epoch,
            "best_metric": self.best_metric,
            "history": self.history,
            "model_config": self.model.__class__.__name__,
        }
        torch.save(checkpoint, self.save_dir / filename)
    
    def load_checkpoint(self, checkpoint_path: str):
        """恢复训练状态"""
        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        if self.scheduler and checkpoint["scheduler_state_dict"]:
            self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        self.current_epoch = checkpoint["epoch"]
        self.best_metric = checkpoint["best_metric"]
        self.history = checkpoint["history"]
        print(f"已从 {checkpoint_path} 恢复训练，当前epoch: {self.current_epoch}")

一个容易被忽略的坑：torch.load的时候一定要指定map_location=self.device。如果你的模型是在GPU上保存的，加载时换了一台只有CPU的机器，或者换了GPU型号，不指定map_location会报错或者产生兼容性问题。这个参数就是为了在加载时把张量映射到当前可用的设备上。

6. 训练器封装之后：回调机制与断点续训的进阶玩法

基础版本有了，再谈进阶。说实话，我之所以觉得用类封装训练流程是必须的，就是因为后续一旦要加功能，过程式的代码会越来越难维持，而类封装能让你以极低的成本不断往里面加东西。

6.1 回调函数机制：让训练过程可插拔

回调机制是我最推荐的一个设计。它的本质是：在训练的关键节点预留钩子，让你不用修改训练器内部代码，就能插入自定义逻辑。

比如你想每训练一个epoch就在TensorBoard上记录loss曲线，你不需要改训练器代码，只需要写一个回调函数，然后把它传进去：

python复制from torch.utils.tensorboard import SummaryWriter


class TensorBoardCallback:
    def __init__(self, log_dir="runs"):
        self.writer = SummaryWriter(log_dir)
    
    def __call__(self, trainer):
        record = trainer.history[-1]  # 最近一次的记录
        self.writer.add_scalar("train_loss", record["train_loss"], record["epoch"])
        self.writer.add_scalar("val_loss", record["val_loss"], record["epoch"])
        self.writer.add_scalar("val_acc", record["val_acc"], record["epoch"])

然后又想加一个早停机制——验证准确率连续N个epoch不提升就停止训练：

python复制class EarlyStopping:
    def __init__(self, patience=3, min_delta=1e-4):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_score = None
    
    def __call__(self, trainer):
        current_score = trainer.history[-1]["val_acc"]
        if self.best_score is None:
            self.best_score = current_score
        elif current_score < self.best_score + self.min_delta:
            self.counter += 1
            if self.counter >= self.patience:
                print(f"早停触发：验证准确率连续{self.patience}个epoch未提升")
                # 注意这里需要和训练器约定一个停止机制
                trainer.stop_training = True
        else:
            self.best_score = current_score
            self.counter = 0

你会发现这个回调本身也用到了__call__——又是一个可调用对象。

在训练器里配合早停，需要在__call__的主循环加一行判断：

python复制        for epoch in range(self.current_epoch, target_epochs):
            self.current_epoch = epoch + 1
            
            # ... 训练和验证 ...
            
            for callback in self.callbacks:
                callback(self)
            
            if getattr(self, "stop_training", False):
                print("训练提前终止")
                break

这个设计模式非常像PyTorch Lightning的Callbacks机制，也是借鉴了Keras的回调设计。虽然我们这里用几十行代码实现了简化版，但核心思想是一样的——训练器负责"训练"，回调负责"观察和干预"，两者通过训练器对象上公开的状态（history、best_metric、current_epoch）来通信。

6.2 断点续训与多实验管理

配合checkpoint，断点续训的完整流程是这样：

python复制# 第一次训练，训练10个epoch
trainer = GPUTrainer(...)
trainer(10)

# 某天发现效果不够好，想再训练5个epoch
trainer = GPUTrainer(...)
trainer.load_checkpoint("checkpoints/best_model.pth")
trainer(5)  # 继续训练5个epoch，总共15个epoch

这里有个细节点：因为初始化的时候current_epoch是从checkpoint里恢复的，所以在__call__的循环里，是从self.current_epoch开始而不是从0开始，保证不会重复计算前面的epoch。

多实验管理方面，我习惯用一个简单的实验配置字典，每次实验用一个唯一的experiment_id命名文件夹：

python复制experiment_config = {
    "experiment_id": "resnet50_bs128_lr1e-3_aug_v2",
    "model": resnet50,
    "batch_size": 128,
    "lr": 1e-3,
    "epochs": 50,
    "save_dir": f"checkpoints/resnet50_bs128_lr1e-3_aug_v2",
    "seed": 42,
}

然后把配置存成JSON，和checkpoint放在一起。这样每个实验的配置、代码版本、checkpoint、日志都在一起，复现的时候一目了然。

6.3 混合精度训练的实际收益

在GPUTrainer里我加了use_amp参数，这是NVIDIA AMP（Automatic Mixed Precision）的封装。实测下来，在支持Tensor Core的显卡上（RTX 20系列及以上），混合精度训练能带来约1.5~3倍的加速，显存占用也几乎减半。

AMP的原理是用FP16（半精度）做前向传播和反向传播，但用FP32（单精度）做参数更新。因为FP16的数值范围小，直接训练容易梯度下溢，所以PyTorch用GradScaler来自动缩放梯度，更新参数前再缩放回去。这些细节autocast和GradScaler都帮你处理了，你只需要：

1. 在前向传播和loss计算外层套with autocast():；
2. 用scaler.scale(loss).backward()替代loss.backward()；
3. 用scaler.step(optimizer)替代optimizer.step()；
4. 用scaler.update()更新缩放因子。

如果在深度学习框架里看到torch.cuda.amp，记住这套固定范式，直接用就行。

有个小坑：如果你的模型里有BatchNorm层，混合精度训练时BatchNorm的统计量是在FP16下计算的，某些极端情况下会导致精度下降。一般训练CNN时问题不大，但是如果训练过程中发现验证准确率上不去，可以把BatchNorm层强制设置成FP32，或者干脆use_amp=False对比一下。

6.4 踩坑实录：显存溢出、NVML初始化失败和训练结果难复现

最后再分享几个GPU训练中非常高频的报错。

第一个：显存溢出（CUDA out of memory）。

text复制RuntimeError: CUDA out of memory. Tried to allocate 512.00 MiB 
(GPU 0 has 8.00 GiB total capacity; 7.50 GiB already allocated; ...)

这个报错最常见的三个原因：batch size太大、输入数据没有释放、其他进程占了显存。排查顺序是：先用nvidia-smi看显存是不是被其他进程占了；然后减小batch size或者用梯度累积代替；最后检查代码里是不是把大张量存储在变量里一直没释放。有时torch.cuda.empty_cache()能缓解碎片化问题，但那是治标不治本。

第二个：WSL下的NVML初始化失败。

text复制failed to initialize NVML: GPU access blocked by the operating system

这个我记忆犹新。在WSL里跑GPU训练时经常遇到，一般是指驱动版本和WSL内核版本不匹配。解决方案是：在Windows宿主机上安装最新版的NVIDIA驱动，并且确保WSL更新到WSL2。WSL2里不需要单独装Linux驱动，因为它是通过Windows的驱动桥接访问GPU的。如果装了Linux下的NVIDIA驱动反而可能造成冲突，这个坑很容易误导人。

第三个：训练结果无法复现。

这不算报错，但比报错更折磨人。你跑了两遍完全相同的代码，结果每次的loss曲线和准确率都不一样。原因通常是三个：

1. 没有固定随机种子；
2. cuDNN的自动调优导致算法选择有随机性；
3. 多线程DataLoader导致的数据顺序不稳定。

最基础的复现配方：

python复制import random
import numpy as np
import torch

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

注意cudnn.deterministic=True会牺牲一点性能，换取完全可复现的结果。我的习惯是：调试阶段开着deterministic，方便排查问题；正式训练的时候关掉，换取最快的速度。

7. 最后的几点体会

写到这里，这一天的笔记内容差不多讲完了。总结一下我今天学到和踩过的东西。

__call__的方法设计，我从一开始觉得"这不过是个语法糖"，到现在真正在训练器里用出来，感觉完全不同。它最大的价值不是省了那一个方法名的打字量，而是把对象的交互方式统一成"调用"这个抽象动作——对一个训练器来说，你不需要记住它的API叫fit还是train还是run，你只需要trainer()，跟调用一个函数一样自然。这对代码的可读性和可维护性都是巨大的提升。

而GPU训练这部分，其实环境配置的坑远比模型本身的坑多。驱动、CUDA、框架三者之间的版本匹配关系，如果你不理解"框架自带CUDA runtime"这个核心事实，很容易陷入反复重装却始终不生效的泥潭。我见过太多人在环境上浪费一整天，最后发现只是装错了PyTorch的wheel包。

如果你正好也在写自己的深度学习训练脚本，我强烈建议你花一个下午的时间，把你的train()大函数重构成一个训练器类，加上__call__方法。不用一次性做完所有功能，先做一个只支持普通精度训练、不保存checkpoint的极简版，跑通之后再慢慢加。你会发现，代码结构变清晰之后，你调参、加实验、排错的速度会快很多。

我自己的下一步计划是给这个训练器加上分布式数据并行（DDP）的支持，因为单卡训练在模型变大之后实在不够用。做法也简单：在多卡环境下，用torch.distributed初始化进程组，然后把模型用torch.nn.parallel.DistributedDataParallel包一层，数据加载器换成DistributedSampler。核心的训练逻辑不用动，仍然是__call__这一套。这个等我自己踩过一圈坑之后，再单独写一篇记录下来。