騰訊發布推理框架TNN全平台版本，同時支援移動端、桌面端和服務端

2021 年 4 月 28 日
AI
人臉識別

TNN是騰訊開源的新一代跨平台深度學習推理框架，也是騰訊深度學習與加速Oteam雲帆的開源協同成果，由騰訊優圖實驗室主導，騰訊光影研究室、騰訊雲架構平台部、騰訊數據平台部等團隊一起協同開發。在經過4個多個月的迭代完善後，TNN新版本v0.3版本正式發布，是首個同時支援移動端、桌面端、服務端的全平台開源版本。TNN新版本在通用性、易用性、性能方面進一步獲得提升。

TNN地址：

//github.com/Tencent/TNN

通用性

在保證模型統一、介面統一的前提下，依託於硬體廠商提供的加速框架基礎運算元支援，以及手寫kernel優化的方式，對移動端、桌面端和服務端提供了多種不同加速選擇，實現了對常用CV、NLP模型的優化適配。

硬體平台支援

TNN通過集成OpenVINO和TensorRT的方式新增了對服務端X86和NVIDIA硬體的支援，既能快速獲取硬體廠商的最新優化成果，又能基於業務模型結構特點添加自定義實現達到性能極致。同時考慮到桌面端應用對安裝包大小的限制，TNN通過JIT和手工優化的方式實現了輕量級的X86後端，整體庫大小僅為5MB左右。

模型運算元支援

TNN新版本在CV類模型的支援上擴展了對3D-CNN、LSTM、BERT等模型結構的支援，總運算元數從88個增加到107個，新增運算元包括LSTM、GridSample、Histogram、OneHot、BitShift、Gather、ScatterND、LayerNorm、GroupNorm、GELU、SoftSign、Erf等。

易用性

動態維度和預處理支援

TNN之前版本主要支援CV類模型，網路輸入基本都是NCHW4個維度且每個維度上的值基本不變。而NLP場景下同一個網路會有0維到6維的情況，且每個維度上的值根據輸入而變化。為此TNN新增了輸入維度配置介面，在模型運算元、硬體、系統支援等層面做了大量補充和完善。

API介面對Mat相關介面做了一定擴充，包括拷貝填充功能(CopyMakeBorder)，方便SDK開發者進行網路預處理後處理加速。目前TNN已支援裁剪(Crop)、縮放(Resize)、顏色空間轉換(CvtColor)、仿射變換(WarpAffine)和拷貝填充(CopyMakeBorder)等常見的預處理後處理操作。

運行時常量摺疊

onnx模型導出模型時會產生很多粘合劑類的運算元用於計算常量和數據Shape的資訊，TNN實現了ConstFolder常量摺疊功能來簡化模型結構提升模型運行性能。相比開源社區工具onnx-simplifier，ConstFolder增加了對以ATen形式輸出的運算元的支援，同時支援運行時常量摺疊以支援模型變維的需求。TNN運行時將變維計算部分的運算元單獨抽取出來用NAIVE（純C++）執行，以減輕各個硬體device（ARM、Metal、OpenCL）的運算元實現壓力。

示例展示

在騰訊光影和騰訊光流團隊的支援下，頭髮染色和人體姿態兩個TNN示例在TNN中間迭代的小版本中已經發布，並展現出不錯的演算法效果。這次配合新版本發布，我們添加了移動端中文OCR示例以及桌面端/後台端BERT閱讀理解示例。

中文OCR示例採用chineseocr lite模型，展示了如何通過文本框位置檢測+文本框角度檢測+文字識別3個模型串聯來進行中文文字識別；

BERT 閱讀理解示例採用BERT-Squad10模型，展示了如何通過預先輸入上下文以及辭彙表來實現一個簡單的問答系統。下圖依次為頭髮染色、人體姿態、中文OCR、BERT閱讀理解效果展示。

性能優化

移動端性能優化

Arm性能優化:

armv8.2優化：fp16向量指令優化，相比於fp32預期性能翻倍，除了和大多數開源框架一樣支援arm64之外，針對arm32架構也實現了fp16指令優化，讓64位和32位APP都能發揮硬體fp16向量加速的能力；

int8優化：針對常用的運算元block組合採取了更激進一些的融合策略，如conv+add+activation，能有效的減少量化反量化的開銷以及記憶體讀寫，並且經內部業務驗證，在提高性能的同時並不會造成精度的下降

OpenCL性能優化：

核心卷積優化：

a. 訪存性能優化: Channel Blocking優化、以及局部記憶體(local memory)優化提升訪存性能，實現工作組內的數據共享；

b. 計算性能優化: winograd演算法優化3×3卷積，定址計算優化，相鄰計算網格偏移量共用向量暫存器，降低fp32計算單元壓力；

工作組尺寸優化: 優化計算策略，並通過Auto-Tuning挑選最優工作組尺寸；

預處理/後處理優化: 使用buffer做參數快取，降低GPU拷貝開銷。

桌面端/服務端性能優化

TNN服務端通過集成OpenVINO和TensorRT的方式新增了對服務端X86和NVIDIA硬體的支援，既能快速獲取硬體廠商的最新優化成果，又能基於業務模型結構特點添加自定義實現達到性能極致。與業界服務端統一框架onnxruntime性能最好版本相比，TNN當前在CV類模型有一定優勢，而onnxruntime在NLP類模型有一定優勢。TNN剛開始支援NLP模型，後續會在這塊持續優化。

TNN桌面端為了兼顧高性能和硬體兼容性，同時考慮應用App 對安裝包大小的限制，通過JIT和手工優化的方式實現了輕量級的X86後端，支援SSE41、SSE42、AVX、AVX2、FMA等指令集。相比onnxruntime服務端庫80MB，TNN桌面端整體庫大小僅為5MB左右，而性能差距在20%以內。