作者 | OpenMMLab（已授權） 

來源 | https://zhuanlan.zhihu.com/p/396781295 

編輯 | 極市平臺

導讀

模型轉換是模型部署的重要環(huán)節(jié)之一，本文會從深度學習訓練框架的角度出發(fā)，講一講作者本人對模型轉換的理解。

1 模型轉換的意義

模型轉換是為了模型能在不同框架間流轉。

在實際應用時，模型轉換幾乎都用于工業(yè)部署，負責模型從訓練框架到部署側推理框架的連接。這是因為隨著深度學習應用和技術的演進，訓練框架和推理框架的職能已經逐漸分化。

分布式、自動求導、混合精度……訓練框架往往圍繞著易用性，面向設計算法的研究員，以研究員能更快地生產高性能模型為目標。

硬件指令集、預編譯優(yōu)化、量化算法……推理框架往往圍繞著硬件平臺的極致優(yōu)化加速，面向工業(yè)落地，以模型能更快執(zhí)行為目標。

由于職能和側重點不同，沒有一個深度學習框架能面面俱到，完全一統(tǒng)訓練側和推理側，而模型在各個框架內部的表示方式又千差萬別，所以模型轉換就被廣泛需要了。

2 模型轉換的技術細節(jié)

2.1 計算圖生成

目前使用廣泛的訓練框架 PyTorch，以及商湯自研的訓練框架 SenseParrots 使用的都是動態(tài)圖，這是由于動態(tài)圖的表達形式更易于用戶快速實現(xiàn)并迭代算法。動態(tài)圖框架會逐條解釋，逐條執(zhí)行模型代碼來運行模型，而計算圖生成是的本質是把動態(tài)圖模型靜態(tài)表達出來。PyTorch 的torchscript、ONNX、fx 模塊都是基于模型靜態(tài)表達來開發(fā)的。目前常見的建立模型靜態(tài)表達的方法有以下三種：

代碼語義分析：通過分析用戶代碼來解析模型結構，建立模型靜態(tài)表達。

模型對象分析：通過模型對象中包含的成員變量，來確定模型算子組成，建立模型靜態(tài)表達。

模型運行追蹤：運行模型并記錄過程中的算子信息、數據流動，建立模型靜態(tài)表達。

上面這三種方法在適用范圍、靜態(tài)抽象能力等方面各有優(yōu)劣。目前訓練框架都主要使用模型運行追蹤的方式來生成計算圖：在模型inference 的過程中，框架會記錄執(zhí)行算子的類型、輸入輸出、超參、參數等算子信息，最后把 inference 過程中得到的算子節(jié)點信息和模型信息結合得到最終的靜態(tài)計算圖。

2.2 計算圖中的自定義算子

很多時候，用戶的一段代碼可能涉及非框架底層的計算，例如下面這段代碼，涉及外部庫的計算，訓練框架自身是無法追蹤記錄到的。

這個時候我們可以把這部分代碼作為一個自定義算子，由用戶定義這個算子在計算圖中作為一個節(jié)點所記錄的信息。實際實現(xiàn)時，這些計算會被寫到一個 Function 或者 Module 中，然后用戶在 Function 或者 Module 中定義這個計算對應的計算節(jié)點的信息表達，這樣每次調用這個定義好的 Function 或者 Module，就能對應在計算圖中記錄相應的算子信息。

當然還有很多其他場景會產生這種需要，例如你的幾個計算組成了一個常見的函數，可以有更高層的表達，這個時候也可以使用自定義算子來簡化計算圖的表達。

2.3 目標格式（caffe/ONNX）

模型轉換往往將模型轉換到一種中間格式，再由推理框架讀取中間格式。

目前主流的中間格式有 caffe 和 ONNX（Open Neural Network Exchange），兩者底層都是基于 protobuf （Google 開發(fā)的跨平臺協(xié)議數據交換格式工具庫）實現(xiàn)的。

caffe 原本是一個經典的深度學習框架，不過由于出現(xiàn)較早且不再維護，已經少有人用它做訓練和推理了。但是它的模型表達方式卻保留了下來，作為中間格式在工業(yè)界被廣泛使用。

ONNX 是各大 AI 公司牽頭共同開發(fā)的一個中間表達格式，用于模型格式交換，目前在社區(qū)非?；钴S，處于不斷更新完善的階段。

由于 caffe 出現(xiàn)較早，在使用上對硬件部署側比較友好（原生算子列表在推理側容易實現(xiàn)，而且 caffe 使用 caffe.proto 作為模型格式數據結構的定義，能實現(xiàn)中心化、多對一），目前很多推理側硬件廠商依然使用 caffe，很多端到端的業(yè)務解決方案，也喜歡使用 caffe。

而 ONNX 有豐富的表達能力、擴展性和活躍的社區(qū)，深受訓練側開發(fā)者、第三方工具開發(fā)者的喜愛， PyTorch 早已將 ONNX 作為官方導出格式進行支持，而 TensorFlow 也非官方地支持 ONNX。

2.4 計算圖轉換到目標格式

計算圖轉換到目標格式就是去解析靜態(tài)計算圖，根據計算圖的定義和目標格式的定義，去做轉換和對齊。這里的主要的工作就是通用的優(yōu)化和轉換，以及大量 corner case 的處理，相信看過 PyTorch 的 ONNX 導出源碼，或者自己做過相關工作的人都深有體會。

2.4.1 計算圖轉換到 caffe

一般支持 caffe 的推理框架都是在原生 caffe 的基礎上自己額外定義了一些算子（有的還會修改一些原生 caffe）的算子，這些改動都能體現(xiàn)在caffe.proto上：例如下圖所示例子就是Mean（PartialMean）在 caffe.proto 中的定義。使用這樣的 proto 文件和 protobuf，才能生成帶 Mean 算子的 caffe 格式模型（.prototxt, .caffemodel）。

這是一個以推理框架為中心的生態(tài)，不同的推理框架提供不同的 caffe.proto，就可以形成各自的算子定義和約束，平時我們把推理框架自己定義的 caffe 格式稱為 caffe 后端。

計算圖轉換到 caffe，就是將計算圖的算子進行分發(fā)映射，轉換到不同的 caffe 后端，計算圖的算子和 caffe 中的算子可能存在一對多、多對一的映射關系。我們遍歷現(xiàn)有的計算圖算子列表，能夠很自然地去處理一對多的轉換映射，而多對一的映射關系就需要針對每個 caffe 后端配置各自的計算圖優(yōu)化pass去預處理計算圖。

2.4.2 計算圖轉換到 ONNX

ONNX 官方定義了算子集 opset，并且隨著 ONNX 的演進，在寫下這篇文章的時候，版本已經迭代到了 opset15。opset 版本的迭代伴隨著算子支持列表和算子表達形式的改動，因此針對不同的 opset 也需要有多后端 ONNX 的支持。另一方面，對于在 opset 之外的算子，用戶需要自己注冊定義算子在 ONNX 的表達信息（輸入、輸出、超參等）。

另一方面，推理框架對于 ONNX 官方 opset 往往也不是完全支持，會有自己的一些取舍。所以對于 ONNX 模型，往往需要用相關的 simplifier 進行模型預處理優(yōu)化，圍繞這一方面模型轉換或者部署框架的工程側也有不少的相關工作。

2.4.3 onnxruntime 和 caffe 的推理能力

和五花八門的芯片等端側硬件相比，x86 和 CUDA 平臺是普及率最高的平臺，因此如果是出于部署測試、轉換精度確認、量化等需要，一個能夠在 x86 或者 CUDA 平臺運行的 runtime 是非常必要的。

對此，支持 ONNX 格式的部署框架一般會基于 onnxruntime（微軟出品的一個具有 ONNX 執(zhí)行能力的框架）進行擴展，支持 caffe 格式的部署框架一般會基于原生 caffe 進行擴展。通過 onnxruntime 和 caffe 的推理運行能力，來提供在 x86 或者 CUDA 平臺上和硬件平臺相同算子表達層次的運行能力。

當然還有一些生態(tài)較好的部署框架，他們自己提供算子表達能力和計算精度與硬件一致的 x86 或 CUDA 平臺的模擬器。

2.5 端到端的模型轉換

還有一些模型轉換是直接從框架到框架對接一步到位的，相比使用中間格式的方案非常定制化。

例如由英偉達官方出品的 CUDA 平臺的部署框架 TensorRT，支持用戶編寫轉換代碼，直接從 PyTorch 轉換到 TensorRT。

這種端到端的模型轉換，是一種拋棄了中間格式的中心化轉換方法，省去了很多麻煩，往往在整個平臺完全自研自主使用，或者業(yè)務構成本身比較單一（解決方案的訓練框架和部署框架完全確定）等實際情況下落地使用。

3 總結

模型轉換是一個由現(xiàn)有的深度學習技術格局和業(yè)務需求衍生出的工程方向，作者這里只是從訓練框架部署的角度介紹了一下自己的意見，相信很多來自其他方向或者接觸其他業(yè)務的人會有著自己的實踐和理解，也非常歡迎大家積極分享交流。

本文僅做學術分享，如有侵權，請聯(lián)系刪文。