MMDeploy项目：如何为后端算子添加单元测试

MMDeploy项目：如何为后端算子添加单元测试

前言

在深度学习模型部署过程中，自定义算子的开发与测试是确保模型在不同推理后端上正确运行的关键环节。本文将详细介绍在MMDeploy项目中如何为自定义后端算子添加单元测试，帮助开发者构建更加健壮的模型部署流程。

准备工作

在开始编写测试之前，需要确保：

1. 已完成自定义算子的开发并放置在backend_ops目录下
2. 已重新编译相关后端以包含新添加的算子

测试框架概述

MMDeploy提供了一个统一的测试框架，支持多种后端（如TensorRT、ONNX Runtime、NCNN等）的算子测试。测试文件位于tests/test_ops/test_ops.py中。

测试用例编写步骤

1. 创建测试函数

测试函数的基本模板如下：

@pytest.mark.parametrize('backend', [TEST_TENSORRT, TEST_ONNXRT])
@pytest.mark.parametrize('op_param1,op_param2', [(val1, val2), (val3, val4)])
def test_custom_op(backend, op_param1, op_param2, input_list=None, save_dir=None):
    # 测试代码

1.1 后端测试类配置

MMDeploy提供了多种后端测试类：

• TestOnnxRTExporter: ONNX Runtime后端测试
• TestTensorRTExporter: TensorRT后端测试
• TestNCNNExporter: NCNN后端测试

通过@pytest.mark.parametrize可以指定需要测试的后端类型。

1.2 算子参数设置

为算子设置不同的参数组合进行测试，例如对于ROI Align算子：

@pytest.mark.parametrize('pool_h,pool_w,spatial_scale,sampling_ratio',
                         [(2, 2, 1.0, 2), (4, 4, 2.0, 4)])

2. 输入数据初始化

根据算子需求初始化输入张量，通常包括：

• 特征图输入
• 算子特定参数（如ROI坐标）
• 其他辅助输入

input = torch.rand(1, 1, 16, 16, dtype=torch.float32)
single_roi = torch.tensor([[0, 0, 0, 4, 4]], dtype=torch.float32)

3. 构建测试模型

有两种方式构建包含自定义算子的测试模型：

3.1 Torch模型方式

from mmcv.ops import roi_align

def wrapped_function(torch_input, torch_rois):
    return roi_align(torch_input, torch_rois, (pool_w, pool_h),
                     spatial_scale, sampling_ratio, 'avg', True)

wrapped_model = WrapFunction(wrapped_function).eval()

3.2 ONNX模型方式

直接使用ONNX API构建模型，适用于更复杂的算子测试场景。

4. 执行后端测试

调用后端测试类的run_and_validate方法进行测试：

with RewriterContext(cfg={}, backend=backend.backend_name, opset=11):
    backend.run_and_validate(
        wrapped_model, [input, single_roi],
        'roi_align',
        input_names=['input', 'rois'],
        output_names=['roi_feat'],
        save_dir=save_dir)

run_and_validate参数详解

| 参数名 | 类型 | 说明 | |--------|------|------| | model | torch/nn.Module | 待测试的模型 | | input_list | list | 测试输入数据列表 | | model_name | str | 模型名称标识 | | tolerate_small_mismatch | bool | 是否允许微小数值误差 | | do_constant_folding | bool | 是否启用常量折叠优化 | | dynamic_axes | dict | 动态维度配置 | | output_names | list | 输出节点名称 | | input_names | list | 输入节点名称 | | expected_result | tensor | 预期结果用于验证 | | save_dir | str | 输出文件保存目录 |

测试执行方法

完成测试编写后，可通过以下命令执行特定算子的测试：

pytest tests/test_ops/test_ops.py::test_custom_op

最佳实践建议

1. 参数覆盖：为算子设计多种参数组合，覆盖典型和边界情况
2. 输入多样性：测试不同形状和数值范围的输入数据
3. 结果验证：对于确定性算子，提供expected_result进行精确验证
4. 多后端测试：尽可能在多个后端上进行测试
5. 错误处理：添加对非法输入的测试用例

结语

通过本文介绍的方法，开发者可以为MMDeploy中的自定义算子构建全面的测试用例，确保算子在各种后端上的正确性和稳定性。良好的测试实践是高质量模型部署的重要保障。