机器学习代码如何写单元测试？

一、先给出结论和背景

• 测确定性部分，不断言精度数值：模型训练本身带随机性，直接断言「准确率 > 0.9」既不稳定也不是单元测试该做的事。单元测试应锁定可确定验证的环节。
• 数据与特征：- 归一化、分桶、缺失值填充等转换：给定输入，断言输出精确等于预期。
• 边界用例：空输入、单样本、含异常值。

二、拆开关键步骤和判断点

• 形状与数值不变量：- 前向输出形状是否符合 (batch, num_classes)。
• 概率输出范围在 [0,1]、softmax 每行和为 1、无 NaN/Inf。
• loss 与梯度：- 用已知输入手算 loss 对比；调用一次 backward 后断言关键参数梯度非 0、量级合理。
  示例代码可以这样理解：
  def test_forward_shape():
  torch.manual_seed(0)
  x = torch.randn(4, 10)
  out = model(x)
  assert out.shape == (4, 3)
  assert torch.allclose(out.softmax(-1).sum(-1), torch.ones(4))

三、补上落地边界和取舍

• Sanity check：让模型在一小批样本上过拟合，loss 应能降到接近 0；降不下去说明管线有 bug。