第三章 精度对齐工具链介绍

做算子精度对齐，说白了就是拿两套计算结果做对比。一套是参考实现（比如 PyTorch 的 CPU 结果），一套是你的硬件实现。但问题来了——怎么比？用什么工具比？

我个人习惯把精度对齐工具分成三类：框架自带工具、第三方对比工具、自定义脚本。今天咱们重点聊聊前两类，也就是 PyTorch、TensorFlow、ONNX Runtime 这些主流框架里自带的精度对齐能力。

3.1 PyTorch 的精度对齐工具

PyTorch 是我用得最多的框架。它的精度对齐工具其实藏在 torch.testing 和 torch.autograd 里。

3.1.1 torch.testing.assert_close

这个函数是 PyTorch 1.9 之后引入的，用来替代老旧的 assert_allclose。用法很简单：

import torch
from torch.testing import assert_close

# 假设这是你的硬件输出
hardware_output = torch.tensor([1.0001, 2.0002, 3.0003])
# 这是 PyTorch 参考输出
reference_output = torch.tensor([1.0, 2.0, 3.0])

# 默认 rtol=1e-5, atol=1e-8
assert_close(hardware_output, reference_output)

嗯，这里要注意：assert_close 默认的容差对于某些算子来说太严格了。我在项目中遇到过，一个简单的 softmax 算子，因为浮点累加顺序不同，误差到了 1e-6 级别。默认的 1e-5 相对误差直接报错。

3.1.2 torch.autograd.gradcheck

这个工具是用来检查梯度计算是否正确的。说白了，它用数值微分去验证你的解析梯度。

from torch.autograd import gradcheck

# 自定义算子
class MyOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        return input * 2
    
    @staticmethod
    def backward(ctx, grad_output):
        return grad_output * 2

input = torch.randn(3, 3, dtype=torch.double, requires_grad=True)
test = gradcheck(MyOp.apply, input, eps=1e-6, atol=1e-4)
print(f"梯度检查通过: {test}")

为什么用 double 类型？因为数值微分本身有截断误差，单精度下容易误判。我曾经踩过这个坑，用 float32 做 gradcheck，死活过不了，换成 double 就一次通过。

3.2 TensorFlow 的精度对齐工具

TensorFlow 这边，我主要用 tf.test.TestCase 和 tf.debugging 模块。

3.2.1 tf.test.TestCase.assertAllClose

这是 TensorFlow 1.x 时代就有的工具，现在依然好用：

import tensorflow as tf

class MyPrecisionTest(tf.test.TestCase):
    def test_conv2d_precision(self):
        # 硬件输出
        hw_result = tf.constant([[[[1.0001]]]], dtype=tf.float32)
        # 参考输出
        ref_result = tf.constant([[[[1.0]]]], dtype=tf.float32)
        
        self.assertAllClose(hw_result, ref_result, rtol=1e-5, atol=1e-5)

说实话，TensorFlow 的测试框架比 PyTorch 更完善。它支持 subTest，可以批量测试不同输入。我习惯把不同 shape、不同 dtype 的测试用例写在一起：

def test_multiple_shapes(self):
    shapes = [(1, 3, 3, 1), (2, 5, 5, 3), (4, 7, 7, 8)]
    for shape in shapes:
        with self.subTest(shape=shape):
            # 测试逻辑
            pass

3.2.2 tf.debugging.assert_near

这个函数和 assertAllClose 类似，但更灵活。它可以在计算图里执行，适合在 eager 模式下调试：

hw = tf.constant([1.0001, 2.0002])
ref = tf.constant([1.0, 2.0])
tf.debugging.assert_near(hw, ref, rtol=1e-5, atol=1e-5)

3.3 ONNX Runtime 的精度对齐工具

ONNX Runtime 作为中间表示，它的精度对齐工具有点特殊。你没法直接拿 ONNX 算子和硬件算子比，因为 ONNX 本身是个中间格式。

3.3.1 使用 ONNX Runtime 做参考

我通常的做法是：把 PyTorch 模型导出成 ONNX，然后用 ONNX Runtime 跑一遍，把结果作为参考值。

import onnxruntime as ort
import numpy as np

# 加载 ONNX 模型
sess = ort.InferenceSession("model.onnx")

# 准备输入
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 运行 ONNX Runtime
outputs = sess.run(None, {"input": input_data})

# 这个 outputs 就是你的参考值
print(outputs[0])

为什么要多此一举？因为 PyTorch 的算子实现和 ONNX Runtime 的算子实现可能不同。比如某些算子，PyTorch 用了 cudnn，而 ONNX Runtime 用了 mkldnn，结果会有微小差异。

3.3.2 自定义精度比较脚本

ONNX Runtime 没有内置的 assert_close 函数，所以我一般自己写：

def compare_results(hw_result, ref_result, rtol=1e-5, atol=1e-8):
    """比较硬件输出和 ONNX Runtime 参考输出"""
    diff = np.abs(hw_result - ref_result)
    max_diff = np.max(diff)
    mean_diff = np.mean(diff)
    
    # 计算相对误差
    ref_abs = np.abs(ref_result)
    rel_diff = diff / (ref_abs + 1e-10)
    max_rel = np.max(rel_diff)
    
    print(f"最大绝对误差: {max_diff:.2e}")
    print(f"平均绝对误差: {mean_diff:.2e}")
    print(f"最大相对误差: {max_rel:.2e}")
    
    # 判断是否通过
    if max_diff > atol and max_rel > rtol:
        print("❌ 精度不达标")
        return False
    else:
        print("✅ 精度通过")
        return True

3.4 工具链选型与对比

说了这么多，到底该选哪个？我画了张图，帮你理清思路：

3.5 工具链对比表

说实话，工具只是辅助。真正重要的是你理解算子的数值特性。比如某些算子对输入范围敏感，某些算子在边界值附近误差大。这些经验，只能靠一个一个算子去踩坑积累。

我记得有一次，一个 LayerNorm 算子，在 PyTorch 和硬件上结果差了 1e-4。我查了三天，最后发现是 epsilon 参数在分母上的处理顺序不同。PyTorch 先加 epsilon 再开方，硬件先开方再加 epsilon。就这么一点顺序差异，导致了精度问题。

所以，工具链选型固然重要，但更重要的是——你要对算子的数学定义、浮点运算的细节有深刻理解。工具只是帮你更快地发现问题，而解决问题，靠的是你对算子的理解。

好了，这一章就到这里。工具链选型没有银弹，根据你的实际场景选择最合适的就好。