TNN项目实战：MobileNetV2-SSD模型转换与部署全流程解析

TNN项目实战：MobileNetV2-SSD模型转换与部署全流程解析

TNN TNN: developed by Tencent Youtu Lab and Guangying Lab, a uniform deep learning inference framework for mobile、desktop and server. TNN is distinguished by several outstanding features, including its cross-platform capability, high performance, model compression and code pruning. Based on ncnn and Rapidnet, TNN further strengthens the support and performance optimization for mobile devices, and also draws on the advantages of good extensibility and high performance from existed open source efforts. TNN has been deployed in multiple Apps from Tencent, such as Mobile QQ, Weishi, Pitu, etc. Contributions are welcome to work in collaborative with us and make TNN a better framework. 项目地址: https://gitcode.com/gh_mirrors/tn/TNN

一、目标检测与SSD算法简介

目标检测是计算机视觉领域的核心任务之一，它需要同时完成物体定位（在哪里）和分类（是什么）两个任务。SSD（Single Shot MultiBox Detector）是一种经典的单阶段目标检测算法，其核心特点是通过在不同尺度的特征图上进行预测，能够高效地检测各种大小的物体。

MobileNetV2作为轻量级网络架构，与SSD结合形成的MobileNetV2-SSD模型，在保持较高检测精度的同时大幅提升了推理速度，非常适合移动端和嵌入式设备的实时检测场景。

二、模型转换整体流程概述

将TensorFlow训练好的MobileNetV2-SSD模型部署到TNN推理框架，需要经历以下几个关键步骤：

1. 模型下载与清理
2. TensorFlow转ONNX格式
3. ONNX模型结构调整
4. 后处理集成（边界框解码）
5. ONNX转TNN格式
6. TNN推理实现

下面我们将详细解析每个步骤的技术要点和实现方法。

三、模型准备与预处理

3.1 模型下载与结构分析

原始模型采用TensorFlow Object Detection API训练得到，包含以下关键文件：

• frozen_inference_graph.pb：冻结的计算图
• saved_model：SavedModel格式模型
• checkpoint：训练检查点

使用Netron工具可视化模型结构，可以清晰看到输入输出节点：

• 输入节点：Preprocessor/sub
• 输出节点：
  • concat：边界框坐标
  • Postprocessor/convert_scores：类别得分

3.2 模型精简与优化

原始模型包含训练专用的节点和子图，需要精简以提升推理效率。我们使用TensorFlow的strip_unused工具进行子图裁剪：

def optimize_graph(graph):
    gdef = strip_unused_lib.strip_unused(
        input_graph_def=graph.as_graph_def(),
        input_node_names=[input_node],
        output_node_names=[bbox_output_node, class_output_node],
        placeholder_type_enum=dtypes.float32.as_datatype_enum)
    return gdef

此操作会保留从指定输入到输出的最短路径，移除所有无关节点，显著减小模型体积。

四、模型格式转换

4.1 TensorFlow转ONNX

使用tf2onnx工具进行转换，关键参数说明：

python3 -m tf2onnx.convert \
  --graphdef saved_model.pb \
  --inputs "Preprocessor/sub:0[1,300,300,3]" \  # 输入节点及shape
  --inputs-as-nchw "Preprocessor/sub:0" \      # 指定NCHW布局
  --outputs "Postprocessor/convert_scores:0,concat:0" \  # 输出节点
  --output saved_model.onnx \                  # 输出路径
  --opset 11 \                                # ONNX算子集版本
  --fold_const                                # 常量折叠优化

转换过程中需要注意：

1. 输入输出节点名称需与原始模型一致
2. 图像尺寸需与训练配置一致（300x300）
3. 建议使用opset 11以获得最佳兼容性

4.2 ONNX模型结构调整

原始ONNX模型的输入输出节点名称不够直观，我们进行语义化重命名：

def modify_io_name(model, src_name, dst_name, io_type='input'):
    io_list = model.graph.input if io_type == 'input' else model.graph.output
    for io in io_list:
        if io.name == src_name:
            io.name = dst_name
    # 更新所有相关节点的输入输出引用
    # ...

修改后：

• 输入节点：input
• 输出节点：
  • score：类别得分
  • box：边界框坐标

五、后处理集成

5.1 边界框解码原理

SSD模型的输出并非直接可用的坐标，而是基于预设Anchor的偏移量。解码公式如下：

解码后坐标 = (预测偏移量 * 缩放因子) * Anchor尺寸 + Anchor中心坐标

其中：

• 中心坐标(x,y)的缩放因子为0.1
• 宽高(w,h)的缩放因子为0.2，且需做指数运算

5.2 ONNX算子实现解码

我们使用ONNX Python API构建解码子图，主要算子包括：

1. Slice：切分输出Tensor

slice_node = helper.make_node(
    'Slice',
    inputs=['box', 'starts', 'ends', 'axes', 'steps'],
    outputs=['slice_out'],
    name='custom_slice'
)

2. 数学运算：实现解码公式

# 乘法节点
mul_node = helper.make_node(
    'Mul',
    inputs=['slice_out', 'scale_factor'],
    outputs=['mul_out'],
    name='decode_mul'
)

# 指数运算节点
exp_node = helper.make_node(
    'Exp',
    inputs=['mul_out'],
    outputs=['exp_out'],
    name='decode_exp'
)

3. Concat：合并解码结果

concat_node = helper.make_node(
    'Concat',
    inputs=['yx_decoded', 'hw_decoded'],
    outputs=['output'],
    axis=3,
    name='decode_concat'
)

5.3 模型合并

将解码子图与主模型合并：

combined_model = onnx.compose.merge_models(
    ssd_model, 
    decode_model,
    io_map=[("box", "box")]  # 连接SSD输出与解码输入
)

合并后模型结构验证要点：

1. IR版本一致性
2. Opset版本一致性
3. 输入输出连接正确性

六、TNN模型转换与部署

6.1 模型转换

使用TNN转换工具将ONNX转为TNN格式：

python3 converter.py onnx2tnn \
    new_model.onnx \
    -align \          # 校验模型正确性
    -optimize         # 启用优化

转换过程会进行以下操作：

1. 算子兼容性检查
2. 模型结构优化
3. 生成TNN模型文件（.tnnproto + .tnnmodel）

6.2 TNN推理实现

在TNN中实现SSD推理主要分为三个步骤：

1. 输入预处理

MatConvertParam GetConvertParamForInput() {
    MatConvertParam param;
    // 归一化参数：2/255缩放，-1平移
    param.scale = {2.0f/255, 2.0f/255, 2.0f/255, 0.0f};
    param.bias  = {-1.0f, -1.0f, -1.0f, 0.0f};
    return param;
}

2. 推理执行

auto status = instance->Forward();
if (status != TNN_OK) {
    LOGE("Forward failed: %s\n", status.description().c_str());
    return;
}

3. 输出后处理

void ProcessSDKOutput() {
    // 获取输出Mat
    auto scores = output->GetMat("score");
    auto boxes = output->GetMat("output");
    
    // 生成候选框
    std::vector<ObjectInfo> objects;
    GenerateObjects(objects, scores, boxes, 0.75f);
    
    // NMS处理
    std::vector<ObjectInfo> results;
    TNN_NS::NMS(objects, results, 0.25f, TNNHardNMS);
}

6.3 效果展示

部署后模型能够实时检测图像中的多个物体，输出效果如下：

[检测结果]
类别: person, 置信度: 0.98, 坐标: [120, 80, 300, 400]
类别: car, 置信度: 0.95, 坐标: [50, 150, 280, 250]

七、性能优化建议

1. 模型层面：

   • 使用TNN的量化工具进行INT8量化
   • 开启TNN模型压缩选项

2. 推理层面：

   • 使用多线程推理
   • 合理设置输入分辨率平衡精度与速度

3. 后处理优化：

   • 使用快速NMS算法
   • 对输出进行分批次处理

通过本教程，我们完整实现了从TensorFlow模型到TNN推理的整个流程，关键是将复杂的后处理集成到模型中，大大简化了部署难度。这种思路同样适用于其他检测模型的部署，具有很好的参考价值。

TNN TNN: developed by Tencent Youtu Lab and Guangying Lab, a uniform deep learning inference framework for mobile、desktop and server. TNN is distinguished by several outstanding features, including its cross-platform capability, high performance, model compression and code pruning. Based on ncnn and Rapidnet, TNN further strengthens the support and performance optimization for mobile devices, and also draws on the advantages of good extensibility and high performance from existed open source efforts. TNN has been deployed in multiple Apps from Tencent, such as Mobile QQ, Weishi, Pitu, etc. Contributions are welcome to work in collaborative with us and make TNN a better framework. 项目地址: https://gitcode.com/gh_mirrors/tn/TNN