將TensorFlow Lite模型轉換為ONNX

文章轉載于微信公眾號:黎明灰燼
作者:黎明灰燼

簡介

由 Facebook 和 Microsoft 創(chuàng)建的開放格式神經(jīng)網(wǎng)絡交換格式 ONNX，是一種用于表示機器學習模型。

圖1：ONNX愿景

TF2ONNX 將 TensorFlow 模型轉換為 ONNX，從而將 TensorFlow 訓練后的模型引入支持 ONNX 的系統(tǒng)。

然而 TF2ONNX 有一些局限（v1.5.5，即開始開發(fā) TFLite2ONNX 的時候），例如不支持 TensorFlow 2.0 或量化。將/_易變/_的 TensorFlow 模型轉換為 ONNX 的工作量很大。并且，由于量化在深度學習部署中扮演著越來越重要的角色。

另一方面，TFLite 的模型表示相對_穩(wěn)定_，并且由 Google 統(tǒng)一維護的 TensorFlow 和 TFLite 的模型轉換器足夠健壯。該轉換器通過包括批量歸一化折疊和激活函數(shù)融合在內(nèi)的圖轉換簡化了 TensorFlow 模型。還可以處理在 TensorFlow Quantization-aware Training（QAT）期間生成的 FakeQuantWithMinMaxVars 節(jié)點。

此外，盡管某些模型由 TensorFlow 構建，但僅發(fā)布時只有 TFLite 格式的模型，例如 Google MediaPipe 。ONNX 生態(tài)系統(tǒng)無法使用這類模型。

TFLite2ONNX 可以將 TFLite 模型轉換為 ONNX。截至 v0.3，TFLite2ONNX 支持 TensorFlow 2.0（感謝 TFLite 轉換器）和量化。本文介紹了 TFLite2ONNX 為縮小 TFLite 與 ONNX 模型表示之間的語義差異的背景和實現(xiàn)。

數(shù)據(jù)布局語義轉換

最明顯的差距是數(shù)據(jù)布局問題—— TFLite 模型是 NHWC 格式，而 ONNX 是NCHW，在本文中被稱為_布局語義差異_。

問題與 TF2ONNX

TFLite 的數(shù)據(jù)布局格式在文檔或模型表示中均未提及，但在 TFLite 轉換器（TensorFlow 模型需要為 NHWC）和內(nèi)核有隱式協(xié)議。ONNX 則在算子表示和文檔（由算子表示生成）中明確聲明它使用NCHW。

圖2：TF2ONNX 的數(shù)據(jù)布局處理—— MobileNetV2 示例

TF2ONNX 將_內(nèi)部_算子和張量轉換為 NCHW 數(shù)據(jù)布局，并允許用戶通過 --inputs-as-nchw 選擇是否需要將圖的輸入和輸出轉換為 NHWC 數(shù)據(jù)布局。默認情況（未指定 NCHW）會插入 Transpose 算子以橋接 NHWC 和 NCHW 子圖。上面的_圖2_是一個使用 TF2ONNX 將 MobileNetV2 TensorFlow模型轉換為 ONNX 的實例。（有關 TF2ONNX 處理數(shù)據(jù)布局的更多描述，請參見 GitHub issue。）

在 TFLite2ONNX 的開發(fā)過程中，我們嘗試了兩種方法：

• _基于轉換的方法_—— v0.1啟用，v0.3 刪除。
• _基于傳播的方法_—— v0.2引入并作為默認方法。

基于轉換的方法

一個關于_布局語義差異_的事實是，某些算子具有隱式數(shù)據(jù)布局，如 Conv；而其他則不是，如 Add。

TFLite2ONNX 的_基于轉置的方法_在算子有_布局語義差異_的地方插入一個_轉置模式_。_轉置模式_是用一個 Transpose 算子將_源布局_（TFLite）和_目標的布局_（ONNX）連接起來。

例如，將 TFLite 模式 ? ? ?→[ ]?Datanhwc?→[Conv] 轉換為 ? ? ?→[ ]→? ? ?→[ ]?Datanhwc?→[Transpose]→?Datanchw?→[Conv]。（在這篇文章中，? ??TensorName? 和 [ ][Operator] 分別表示張量和算子。_圖3_是轉換MobileNetV2 的第一個 Conv 的示例。

圖3：通過 TFLite2ONNX 的基于轉置方法轉換的 ONNX 模型

使用這種方法，我們只需要處理一組有限的算子，例如 Conv 和 Pooling。其他的算子和張量轉換都是平凡的——沒有_布局語義上的差異_。

基于傳播的方法

盡管_基于轉換的方法_可以處理_布局語義差異_，但由于添加了太多的算子和張量（即_轉換模式_），因此生成的 ONNX 模型太大且復雜?；趥鞑サ姆椒梢栽谡麄€圖中傳播_布局語義差異_來解決這個問題。

默認情況下（對于大多數(shù)情況），對于給定的圖，某些張量具有隱式布局語義，例如直接連接到 Conv 的張量，而其他張量則沒有，例如 Abs 和 Add。后著的對布局是透明的，這意味著連接到算子的所有張量都必須具有相同的布局語義或不具有這種語義。

因此，當布局_透明_的算子連接到具有隱式布局張量的算子時，_透明_算子的所有張量都具有與連接這兩個算子的張量相同的布局語義。這便是_傳播_的含義。

例如，在轉換 TFLite 圖（省略了_kernel_和 _bias_） ? ? ?→[ ]→? ? ?→[ ]→? ???Anhwc?→[Conv]→?Bnhwc?→[Abs]→?C?? 到 ONNX 時，張量 ? ? ??Anhwc? 變成 ? ? ??Anchw? ，? ? ??Bnhwc? 變成 ? ? ??Bnchw?。因此 [ ][Abs] 的輸出? ??C? 應該與其輸入 ? ??B? 具有相同的格式。_基于傳播的方法會_將 ? ??B? 的格式傳播給 ? ??C?。因此我們得到 ONNX 圖 ? ? ?→[ ]→? ? ?→[ ]→? ? ??Anchw?→[Conv]→?Bnchw?→[Abs]→?Cnchw?，這其中未引入其他算子或張量。

在布局傳播中，如果張量是 _activations_，則布局變換會置換張量的形狀（即 ONNX 中的 value info），如果是_權重_的數(shù)據(jù)（即 ONNX 中的 initializer），還要轉換數(shù)據(jù)。

在實踐中，算子分為四類（_如圖5所示_）：

• _Implicit_：算子在_布局語義上有分歧_，例如 Conv 。它們是_布局語義差異_的來源。
• _Transparent_：對布局不敏感的算子，例如 Abs。如果任何張量具有_布局語義差異_，則將其傳播到連接到此類算子的所有張量。
• Attribute_：可以想 _Transparent 那樣傳播_布局語義差異_的算子，但需要特殊處理處理敏感屬性，例如 Concat 的 axis 屬性。傳播后需要額外通過以調(diào)整這些屬性。
• _Terminate_：沒有和不能傳播_布局語義差異的_算子，例如 Reshape。傳播在碰到此類算子時停止。

圖5：通過基于傳播的TFLite2ONNX方法生成的ONNX模型的一部分

在整個圖中傳播_布局語義差異_時，對于某個算子：如果它是 Transparent 或 Attribute_，則在其張量之間傳播_布局語義差異_；如果是 _Implicit 或 Terminate_，則終止此方向上的傳播。_圖 5 是用_傳播基礎的方法_從 NASNet TFLite 模型轉換得到的 ONNX 模型的一部分。

顯式布局和廣播

通過_基于傳播的方法_，轉換后的 ONNX 模型可輕松處理_布局語義差異_，即無需引入其他算子或張量。

但是，有時可能存在不兼容的布局?？紤]如下的 Reshape。如果 ? ??A? 被傳播而其他張量沒有，由于用戶可能會假設 ? ??B? 的維度和 ? ??A? 有某種關聯(lián)，那么輸出布局可能是意料之外的。（_基于轉換的方法_沒有問題，因為它的布局在模型級別上是 TFLite 格式的，_布局語義差異_在內(nèi)部用 [ ]→[ ]→[ ][Transpose]→[OP]→[Transpose] 模式處理。） { ?}→? ?→[ ? ]→? ?? ?}→[ ]→? ?{Graph}→?A?→[Reshape]→?B??C?}→[Concat]→?D? 我們引入了_顯式布局_來處理這種情況。用戶可以給 TFLite2ONNX 提供 { : ( , )}{Tensor name:tuple(TFLite layout,ONNX layout)} 映射來描述 TFLite 布局和 ONNX 布局的關聯(lián)。而且，用戶可以靈活地為非 Transparent 的算子定義布局轉換。例如，我們對只有 Add 算子的 TFLite 圖執(zhí)行 NHWC 到 NCHW 布局的轉換。

另一個問題是二元算子的廣播，例如 Add（有關更多信息，請參見此問題）。在下面的例子中， ? ??B? 需要廣播。如果 ? ??A? 從 NHWC 轉換為 NCHW，即 ? (2×5×3×4)??A(2×5×3×4)?，而 ONNX 模型中的 ? ??B? 無法廣播。更麻煩的是，_布局語義轉換_在 ? ??B? 處無法傳播，因為 ? ??A? 和 ? ??B? 具有不同的維度。 { ?}→? (2×3×4×5)?? (4×5)?}→[ ]→? ?{Graph}→?A(2×3×4×5)??B(4×5)?}→[Add]→?C? tflite2onnx 引入 _Reshape 模式_來處理廣播問題。對于像 ? ??B? 這樣的張量，拓展它的維度（插入1）使它們彼此相等，以便傳播和廣播可以正確地工作。傳播前的中間圖示例如下。 { ?}→? (2×3×4×5)?? (4×5)?→[ ? ]→? ′(1×1×4×5)?}→[ ]→? ?{Graph}→?A(2×3×4×5)??B(4×5)?→[Reshape]→?B(1×1×4×5)′?}→[Add]→?C?

量化語義轉換

TensorFlow 很早就提供了生產(chǎn)級的量化支持。通過將量化的 TFLite 模型轉換為 ONNX，我們可以將量化功能引入更多系統(tǒng)。（如果本節(jié)中的一些描述使您感到困惑，可以先閱讀神經(jīng)網(wǎng)絡量化簡介。）

問題與 TF2ONNX

TensorFlow 和 TFLite 提供了許多量化解決方案：規(guī)范，后訓練和量化感知訓練。所有這些技術最后生成量化的 TFLite 模型——大多數(shù)情況下時 uint8 格式。這些模型由 TFLite 運行時中的量化版本算子運行。本文將量化張量的 uint8 數(shù)據(jù)、_scale_、_zero point_ 表示為_量化語義_。

另一方面，ONNX中的量化支持有兩個方面（wiki）：

• 接受低精度整數(shù)張量（uint8 或 int8）的量化算子。
• [QLinearConv](https://link.zhihu.com/?target=https%3A//github.com/onnx/onnx/blob/master/docs/Operators.md%23QLinearConv) 和 [QLinearMatMul](https://link.zhihu.com/?target=https%3A//github.com/onnx/onnx/blob/master/docs/Operators.md%23QLinearMatMul) 產(chǎn)生低精度輸出，類似于 TFLite 的量化版 Conv。
• [ConvInteger](https://link.zhihu.com/?target=https%3A//github.com/onnx/onnx/blob/master/docs/Operators.md%23ConvInteger) 和 [MatMulInteger](https://link.zhihu.com/?target=https%3A//github.com/onnx/onnx/blob/master/docs/Operators.md%23matmulinteger) 生成 int32 輸出，可以將其重新量化為低精度。
• [QuantizeLinear](https://link.zhihu.com/?target=https%3A//github.com/onnx/onnx/blob/master/docs/Operators.md%23QuantizeLinear) 以及分別 [DequantizeLinear](https://link.zhihu.com/?target=https%3A//github.com/onnx/onnx/blob/master/docs/Operators.md%23DequantizeLinear) 將高精度（ float 和 int32 ）與低精度轉換的算子。

TensorFlow 和 ONNX 之間的語義鴻溝很大。

在 TensorFlow 生態(tài)中，由于量化表示是為 TFLite 設計的，TensorFlow 圖量化支持有限。因此，TF2ONNX 不提供量化支持。

使用量化算子

在 TFLite2ONNX 最初的設計中，如果量化的 TFLite 算子具有在 ONNX 中有對應，則將其轉換為量化的 ONNX 算子，如 QLinearConv；否則轉換回浮點算子。

由于只有 Conv 和 MatMul 在 ONNX 中具有量化算子，我們不可能生成端到端的量化 ONNX 模型。因此，在量化的 ONNX 算子兩端需要插入 Quantize 和 Dequantize。 ? ?? ?}→[ ]→? ?→[ ]→? ??Aq??Bq?}→[Addq]→?Cq?→[Convq]→?Fq? 例如，給定上面的 TFLite 圖，其中 q 表示張量或算子被量化，量化和反量化算子被插入 [ ][Conv] 兩端， 并將其他地方的張量和算子轉換回浮點，結果如下所示。 ? ?? ?}→[ ]→? ?→[ ]→? 8?→[ ]→? 8?→[ ]→? ??Afloat??Bfloat?}→[Add]→?Cfloat?→[QuantizeLinear]→?Duint8?→[QLinearConv]→?Euint8?→[DequantizeLinear]→?Ffloat? 對于主要由 Conv 構成的模型，例如 MobileNetV1（我們確實嘗試過轉換），這個問題還不大。但對于大多數(shù)其他模型，Conv 和 MatMul 只占算子總數(shù)的一小部分，這要在 ONNX 模型中插入太多的新算子和張量。

而且，像其他許多深度學習系統(tǒng)一樣，ONNX 張量表示不具有量化語義。也就是說，低精度 uint8 張量就是單純的 uint8 數(shù)據(jù)，就像 numpy 一樣——沒有/_scale/_ 和 zero point 描述。對于轉換回浮點的張量，它們的_量化語義_已經(jīng)丟失——這導致我們無法從量化感知訓練中獲益。

維護量化信息

TFLite2ONNX 不使用量化算子，而是通過插入_量化模式_在 ONNX 模型中維護_量化語義_。 [ ]→? ?→[ ]→?????? ?? /_ ?? ??????→[ ]→? ′ ?→[ ][OP]→?Tf?→[Quantize]→{?Tq??Tzero/_point??Tscale?}→[Dequantize]→?Tf′?→[OP] 具體而言，上面的 ONNX 圖是tflite2onnx 是從 TFLite 圖 [ ]→? ?→[ ][OPq]→?Tq?→[OPq] 生成的。

如果原始的 TFLite 模型具有 O 個算子和 T 個張量，則生成的模型中最多可能有 +2 O+2T 個算子和 3 3T 個張量。盡管這種機制增加了更多的張量，但成功在 ONNX 模型中保留了_比例_和_零點_語義。_圖6_ 是將一個量化的 TFLite Conv 模型轉換為 ONNX 的示例。

圖6：由 TFLite2ONNX 生成的量化 ONNX 模型

運行 ONNX 模型的框架可以決定如何啟用量化的 ONNX 模型?？梢詫⒘炕瘓D轉換回非量化圖，或者使用其量化版本算子優(yōu)化_量化模式_，以獲得更好的性能。

實現(xiàn)

截至 v0.3，TFLite2ONNX 是一個非常簡單的僅包含約 2000 行代碼的軟件包。這些代碼分為幾個部分：每個 TFLite 算子專用的轉換器類；Graph 級別管理的數(shù)據(jù)布局和量化處理；幫助函數(shù)或封裝，例如Tensor，Layout。

截至 v0.3 ，許多卷積神經(jīng)網(wǎng)絡已經(jīng)得到支持（測試分支包含了一部分）。支持大約 20 個 TFLite 算子。有命令行工具和 Python 接口可用。

目前的限制包括：

• 尚未啟用 RNN 或控制流網(wǎng)絡。我們計劃在下一個里程碑中做些嘗試。
• 不支持控制生成的 ONNX 模型的版本控制。為了讓 TFLite2ONNX 盡量簡單，目前的計劃是不支持版本控制，而是將其工作留給 ONNX 版本轉換器。
• 目前許多算子尚不支持。如果發(fā)現(xiàn)尚不支持某些算子，可以發(fā)起新算子支持請求。如果您有興趣幫助支持新算子，可以查看這一文檔。

您可以在https://link.zhihu.com/?target=https%3A//github.com/jackwish/tflite2onnx/issues%3Fq%3Dis%253Aissue%2Blabel%253AStory">https://github.com/jackwish/tflite2onnx/issues?q=is%3Aissue+label%3AStory">帶有 Story 標記的 GitHub 問題中找到更多開發(fā)相關的背景。

推薦閱讀

• 深度學習量化技術科普
• 簡單粗暴的多對象目標跟蹤神器 – DeepSort
• 低比特量化之DoreFa-Net理論與實踐

審核編輯：符乾江