第3章基础测试框架设计

3、基础测试框架设计：测试用例基类设计、参数化测试装饰器、自动设备选择（CPU/GPU）

好，咱们进入正题。测试框架这东西，说白了就是给代码上「保险」。你想想看，SwiGLU 激活函数虽然数学上不复杂，但放到不同设备、不同精度、不同输入形状下跑，坑可不少。我早期做模型训练时，就因为没写测试，一个精度问题排查了整整两天——后来发现是 GPU 上的 float16 截断搞的鬼。

所以这一章，咱们就搭一个轻量但够用的测试框架。核心就三件事：基类设计、参数化测试、自动设备选择。

3.1 测试用例基类设计

先说说基类。我个人的习惯是，所有测试用例都继承一个自定义的基类。为什么？因为这样可以统一管理「设备切换」「精度控制」「结果比对」这些重复逻辑。

来看代码：

import torch
import torch.nn as nn
import unittest

class SwiGLUTestBase(unittest.TestCase):
    """SwiGLU 测试基类，所有测试用例都继承这个类"""
    
    def setUp(self):
        """每个测试用例执行前自动调用"""
        self.device = self._get_device()
        self.dtype = torch.float32  # 默认精度，子类可以覆盖
        self.epsilon = 1e-5         # 数值容忍度
        
    def _get_device(self):
        """自动选择设备，后面会详细讲"""
        return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    def _assert_tensor_equal(self, a, b, msg=None):
        """封装断言，支持不同设备间的比较"""
        a = a.cpu().detach() if a.is_cuda else a.detach()
        b = b.cpu().detach() if b.is_cuda else b.detach()
        self.assertTrue(
            torch.allclose(a, b, atol=self.epsilon),
            msg or f"张量不相等，最大差异: {(a - b).abs().max().item()}"
        )
    
    def _create_input(self, shape, requires_grad=True):
        """统一创建输入张量，省得每个测试都写一遍"""
        return torch.randn(shape, device=self.device, 
                          dtype=self.dtype, requires_grad=requires_grad)

核心设计思路：

setUp() 自动初始化设备和精度
_assert_tensor_equal() 封装了设备迁移和精度比较
_create_input() 统一输入创建，避免重复代码

嗯，这里要注意：setUp() 是 unittest 框架的生命周期方法，每个测试用例执行前都会调用一次。所以设备选择、精度设置这些「前置条件」放这里最合适。

3.2 参数化测试装饰器

接下来是参数化测试。说白了就是「同一个测试逻辑，换不同的参数跑多遍」。比如 SwiGLU 的输入维度、batch size、序列长度，这些组合起来可能有几十种情况。你总不能每个组合写一个测试函数吧？

我个人喜欢用 parameterized 库，轻量又好用。安装命令：

pip install parameterized

来看具体用法：

from parameterized import parameterized

class TestSwiGLUForward(SwiGLUTestBase):
    """测试 SwiGLU 前向传播"""
    
    @parameterized.expand([
        ("batch_1_dim_64",   (1, 64)),    # 小批量
        ("batch_4_dim_128",  (4, 128)),   # 常规尺寸
        ("batch_16_dim_256", (16, 256)),  # 大批量
        ("batch_1_dim_1024", (1, 1024)),  # 高维度
    ])
    def test_forward_shape(self, name, shape):
        """验证输出形状是否正确"""
        x = self._create_input(shape)
        model = SwiGLU(shape[-1]).to(self.device)
        output = model(x)
        
        # SwiGLU 输出维度是输入的一半
        expected_shape = (shape[0], shape[-1] // 2)
        self.assertEqual(output.shape, expected_shape,
                        f"{name}: 输出形状不匹配")
    
    @parameterized.expand([
        ("float32", torch.float32, 1e-5),
        ("float16", torch.float16, 1e-3),  # 半精度容忍度更大
    ])
    def test_forward_dtype(self, name, dtype, epsilon):
        """验证不同精度下的前向传播"""
        self.dtype = dtype
        self.epsilon = epsilon
        x = self._create_input((4, 128))
        model = SwiGLU(128).to(self.device, dtype=dtype)
        
        output = model(x)
        self.assertEqual(output.dtype, dtype,
                        f"{name}: 输出精度不正确")

避坑指南：我曾经在参数化测试中踩过一个坑——测试名称不能重复。如果两个测试用例用了相同的 name 参数，unittest 会直接报错。所以建议用「场景_参数」的命名方式，比如 "batch_4_dim_128"。

你可能会问：为什么不直接用 pytest.mark.parametrize？嗯，我个人觉得 unittest + parameterized 的组合更稳定，尤其是在 CI/CD 流水线中，unittest 的兼容性更好。当然，如果你团队统一用 pytest，那也没问题。

3.3 自动设备选择（CPU/GPU）

最后说说自动设备选择。这个看似简单，但坑不少。我见过不少项目直接在代码里写死 device='cuda'，结果在没 GPU 的机器上直接崩掉。

我的做法是：在基类中统一处理，子类无感知。前面 _get_device() 方法已经展示了基本逻辑，但实际项目中还需要考虑更多：

class DeviceManager:
    """设备管理器，支持自动选择和手动指定"""
    
    def __init__(self, prefer_gpu=True):
        self.prefer_gpu = prefer_gpu
        self._device = None
    
    @property
    def device(self):
        if self._device is None:
            self._device = self._auto_select()
        return self._device
    
    def _auto_select(self):
        """自动选择逻辑"""
        if self.prefer_gpu and torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            print(f"✅ 使用 GPU: {gpu_name}")
            return torch.device('cuda')
        else:
            print("⚠️ 使用 CPU（GPU 不可用或未启用）")
            return torch.device('cpu')
    
    def force_cpu(self):
        """强制使用 CPU，用于调试"""
        self._device = torch.device('cpu')
        print("🔧 强制使用 CPU 模式")
    
    def force_gpu(self, device_id=0):
        """强制使用指定 GPU"""
        if torch.cuda.is_available():
            self._device = torch.device(f'cuda:{device_id}')
            print(f"🔧 强制使用 GPU: {torch.cuda.get_device_name(device_id)}")
        else:
            raise RuntimeError("GPU 不可用，无法强制切换")

重要提醒：自动设备选择不是「有 GPU 就用 GPU」这么简单。我遇到过的情况：

GPU 显存不足，强行分配导致 OOM
多 GPU 环境下，默认选了 cuda:0 但其他卡空闲
MPS（Apple Silicon）设备，需要单独处理

所以建议在测试框架中保留 force_cpu() 和 force_gpu() 接口，方便调试。

3.4 知识体系总览

为了让你更直观地理解这三部分的关系，我画了一张图：

这张图把三者的关系说得很清楚了。基类是骨架，参数化是血肉，设备选择是神经系统——三者缺一不可。

3.5 完整测试示例

最后，我把这三部分整合起来，给你看一个完整的测试用例长什么样：

class TestSwiGLUIntegration(SwiGLUTestBase):
    """SwiGLU 集成测试"""
    
    def setUp(self):
        super().setUp()  # 调用父类初始化
        self.model = SwiGLU(128).to(self.device)
    
    @parameterized.expand([
        ("cpu_fp32", 'cpu', torch.float32, 1e-5),
        ("gpu_fp32", 'cuda', torch.float32, 1e-5),
        ("gpu_fp16", 'cuda', torch.float16, 1e-3),
    ])
    def test_forward_backward(self, name, device, dtype, epsilon):
        """前向 + 反向传播联合测试"""
        # 手动切换设备（模拟不同场景）
        if device == 'cuda' and not torch.cuda.is_available():
            self.skipTest("GPU 不可用")
        
        self.device = torch.device(device)
        self.dtype = dtype
        self.epsilon = epsilon
        
        model = SwiGLU(128).to(self.device, dtype=dtype)
        x = self._create_input((4, 128))
        
        # 前向
        output = model(x)
        loss = output.sum()
        
        # 反向
        loss.backward()
        
        # 验证梯度存在且不为零
        for name, param in model.named_parameters():
            if param.grad is not None:
                self.assertFalse(
                    torch.isnan(param.grad).any(),
                    f"{name}: 梯度出现 NaN"
                )
        
        print(f"✅ {name}: 前向反向测试通过")

个人经验：写测试时，我习惯在每个测试用例末尾加一个 print 语句。虽然 unittest 默认只显示失败信息，但加个打印可以让你在跑大量测试时，快速知道「跑到哪了」。当然，正式 CI 环境建议用日志代替 print。

好了，基础测试框架就搭到这里。你可能会觉得「就这？」，但别小看这三板斧——我见过不少项目，连最基本的设备自动选择都没做好，导致测试用例只能在开发者的本地机器上跑。框架不在复杂，在于实用。

📚 课程目录

3、基础测试框架设计：测试用例基类设计、参数化测试装饰器、自动设备选择（CPU/GPU）

3.1 测试用例基类设计

3.2 参数化测试装饰器

3.3 自动设备选择（CPU/GPU）

3.4 知识体系总览

3.5 完整测试示例