將TensorFlow Lite模型轉(zhuǎn)換為ONNX

文章轉(zhuǎn)載于微信公眾號:黎明灰燼
作者:黎明灰燼

簡介

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

圖1：ONNX愿景

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

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

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

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

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

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

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

問題與 TF2ONNX

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

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

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

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

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

基于轉(zhuǎn)換的方法

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

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

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

圖3：通過 TFLite2ONNX 的基于轉(zhuǎn)置方法轉(zhuǎn)換的 ONNX 模型

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

基于傳播的方法

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

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

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

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

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

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

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

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

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

顯式布局和廣播

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

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

另一個(gè)問題是二元算子的廣播，例如 Add（有關(guān)更多信息，請參見此問題）。在下面的例子中， ? ??B? 需要廣播。如果 ? ??A? 從 NHWC 轉(zhuǎn)換為 NCHW，即 ? (2×5×3×4)??A(2×5×3×4)?，而 ONNX 模型中的 ? ??B? 無法廣播。更麻煩的是，_布局語義轉(zhuǎn)換_在 ? ??B? 處無法傳播，因?yàn)?? ??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?

量化語義轉(zhuǎn)換

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

問題與 TF2ONNX

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

另一方面，ONNX中的量化支持有兩個(gè)方面（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 ）與低精度轉(zhuǎn)換的算子。

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

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

使用量化算子

在 TFLite2ONNX 最初的設(shè)計(jì)中，如果量化的 TFLite 算子具有在 ONNX 中有對應(yīng)，則將其轉(zhuǎn)換為量化的 ONNX 算子，如 QLinearConv；否則轉(zhuǎn)換回浮點(diǎn)算子。

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

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

維護(hù)量化信息

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

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

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

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

實(shí)現(xiàn)

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

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

目前的限制包括：

• 尚未啟用 RNN 或控制流網(wǎng)絡(luò)。我們計(jì)劃在下一個(gè)里程碑中做些嘗試。
• 不支持控制生成的 ONNX 模型的版本控制。為了讓 TFLite2ONNX 盡量簡單，目前的計(jì)劃是不支持版本控制，而是將其工作留給 ONNX 版本轉(zhuǎn)換器。
• 目前許多算子尚不支持。如果發(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 標(biāo)記的 GitHub 問題中找到更多開發(fā)相關(guān)的背景。

推薦閱讀

• 深度學(xué)習(xí)量化技術(shù)科普
• 簡單粗暴的多對象目標(biāo)跟蹤神器 – DeepSort
• 低比特量化之DoreFa-Net理論與實(shí)踐