TNN: 腾讯优图实验室和光影实验室开发的高性能、轻量级深度学习推理框架

TNN简介

TNN是由腾讯优图实验室和光影实验室联合开发的一款高性能、轻量级的深度学习推理框架。它在ncnn和Rapidnet的基础上,进一步加强了对移动设备的支持和性能优化,同时借鉴了业界主流开源框架的优点,具有良好的可扩展性和高性能。TNN已在腾讯多个应用中得到广泛应用,如手机QQ、微视、天天P图等。

TNN具有以下主要特点:

1. 跨平台能力强:支持iOS、Android、嵌入式Linux、Windows、Linux等主流操作系统,兼容ARM CPU、X86 GPU、NPU等硬件平台。

2. 高性能:针对不同架构特点进行了深度优化,在主流硬件平台上性能出色。支持INT8/FP16低精度计算加速。

3. 模型压缩:支持INT8量化、剪枝等模型压缩技术,可大幅降低模型大小和内存占用。

4. 代码裁剪:采用模块化设计,可根据实际需求裁剪代码,降低库体积。

5. 统一的推理接口:提供跨平台统一的模型推理接口,方便开发和部署。

6. 丰富的工具链:提供模型转换、量化、可视化等丰富的工具,方便用户使用。

7. 良好的扩展性:采用工厂模式设计,易于支持新的硬件和加速方案。

TNN架构

TNN的整体架构如下图所示:

TNN主要包括以下几个核心模块:

1. 模型解析与图构建:负责解析不同格式的模型,构建计算图。

2. 图优化:进行算子融合、内存复用等图级优化。

3. 算子实现:包含各硬件平台上高度优化的算子实现。

4. 调度执行:负责算子调度和执行。

5. 硬件适配层:抽象硬件接口,方便支持新硬件。

TNN支持通过ONNX导入TensorFlow、PyTorch、MXNet、Caffe等主流训练框架的模型。目前已支持100多种ONNX算子,覆盖了大部分主流CNN、NLP模型所需的算子。

TNN性能

TNN在主流硬件平台上进行了深度优化,性能表现优异。以下是部分常用模型在TNN上的性能数据:

模型	骁龙855(ms)	麒麟980(ms)
MobileNetV2	5.04	6.11
ResNet50	29.91	39.43
YOLOv3-tiny	15.55	19.93

更多性能数据可参考TNN Benchmark。

TNN应用案例

TNN已在腾讯多个业务中得到广泛应用,包括但不限于:

• 人脸检测与对齐
• 人体姿态估计
• 物体检测
• 图像分割
• OCR文字识别
• 阅读理解

以下展示了部分应用效果:

快速上手

使用TNN非常简单,只需三步即可完成模型部署:

1. 模型转换:将训练好的模型转换为TNN模型。TNN提供了丰富的工具支持各种框架模型的转换。

2. 编译TNN:根据目标平台编译TNN推理引擎。TNN提供了方便的一键编译脚本。

3. 集成使用:在应用中调用TNN API进行推理。TNN提供了丰富的Demo作为参考。

详细的使用教程可参考TNN文档。

TNN技术特色

TNN在实现过程中采用了多项先进技术,主要包括:

1. 计算优化

   • 针对不同架构特点进行深度优化
   • 支持Winograd、Tile-GEMM等高效卷积算法
   • 支持算子融合,减少内存访问和kernel启动开销

2. 低精度加速

   • 支持INT8/FP16模式
   • 支持INT8 Winograd算法
   • 支持混合精度

3. 内存优化

   • 高效内存池实现,可减少90%内存开销
   • 支持跨模型内存复用

4. 模块化设计

   • 抽象隔离模型解析、图优化、硬件适配等模块
   • 采用工厂模式注册构建设备

5. 轻量级设计

   • 移动端动态库仅400KB左右
   • 提供基础图像处理操作

TNN能力

TNN目前已具备以下主要能力:

• 支持100+种常用算子
• 支持主流CNN、NLP等模型
• 支持ARM、GPU、NPU等硬件
• 支持Android、iOS、Linux等系统
• 支持模型量化、剪枝
• 支持性能分析

详细的支持情况可参考TNN文档。

开发者指南

TNN欢迎开发者参与贡献,共同打造业界领先的推理框架。TNN提供了详细的开发者文档:

• 开发基础
• 详细架构
• 新增算子
• 单元测试

总结

TNN作为一款高性能、轻量级的深度学习推理框架,在跨平台、性能优化、模型压缩等方面具有显著优势。它已在腾讯多个业务中得到广泛应用,并持续向更多场景拓展。TNN项目采用开源协作的方式,欢迎开发者参与贡献,共同推动深度学习在端侧的发展与应用。

未来,TNN将继续在以下方向持续优化:

1. 支持更多新硬件和加速方案
2. 提升量化模型精度和性能
3. 支持更多模型结构和算子
4. 提供更丰富的工具链
5. 探索端云协同推理

相信在开源社区的共同努力下,TNN必将成为业界领先的深度学习推理框架,为AI技术的普及应用做出重要贡献。