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

TNN简介

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

TNN具有以下主要特点:

跨平台能力强:支持iOS、Android、嵌入式Linux、Windows、Linux等主流操作系统,兼容ARM CPU、X86 GPU、NPU等硬件平台。
高性能:针对不同架构特点进行了深度优化,在主流硬件平台上性能出色。支持INT8/FP16低精度计算加速。
模型压缩:支持INT8量化、剪枝等模型压缩技术,可大幅降低模型大小和内存占用。
代码裁剪:采用模块化设计,可根据实际需求裁剪代码,降低库体积。
统一的推理接口:提供跨平台统一的模型推理接口,方便开发和部署。
丰富的工具链:提供模型转换、量化、可视化等丰富的工具,方便用户使用。
良好的扩展性:采用工厂模式设计,易于支持新的硬件和加速方案。

TNN架构

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

TNN架构图

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

模型解析与图构建:负责解析不同格式的模型,构建计算图。
图优化:进行算子融合、内存复用等图级优化。
算子实现:包含各硬件平台上高度优化的算子实现。
调度执行:负责算子调度和执行。
硬件适配层:抽象硬件接口,方便支持新硬件。

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文字识别
阅读理解

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

人脸检测

人体姿态估计

中文OCR

快速上手

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

模型转换:将训练好的模型转换为TNN模型。TNN提供了丰富的工具支持各种框架模型的转换。
编译TNN:根据目标平台编译TNN推理引擎。TNN提供了方便的一键编译脚本。
集成使用:在应用中调用TNN API进行推理。TNN提供了丰富的Demo作为参考。

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

TNN技术特色

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

计算优化
- 针对不同架构特点进行深度优化
- 支持Winograd、Tile-GEMM等高效卷积算法
- 支持算子融合,减少内存访问和kernel启动开销
低精度加速
- 支持INT8/FP16模式
- 支持INT8 Winograd算法
- 支持混合精度
内存优化
- 高效内存池实现,可减少90%内存开销
- 支持跨模型内存复用
模块化设计
- 抽象隔离模型解析、图优化、硬件适配等模块
- 采用工厂模式注册构建设备
轻量级设计
- 移动端动态库仅400KB左右
- 提供基础图像处理操作

TNN能力

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

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

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

开发者指南

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

总结

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

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

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

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