作者 | OpenMMLab（已授權(quán)） 

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

編輯 | 極市平臺

導讀

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

1 模型轉(zhuǎn)換的意義

模型轉(zhuǎn)換是為了模型能在不同框架間流轉(zhuǎn)。

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

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

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

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

2 模型轉(zhuǎn)換的技術(shù)細節(jié)

2.1 計算圖生成

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

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

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

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

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

2.2 計算圖中的自定義算子

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

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

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

2.3 目標格式（caffe/ONNX）

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

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

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

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

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

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

2.4 計算圖轉(zhuǎn)換到目標格式

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

2.4.1 計算圖轉(zhuǎn)換到 caffe

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

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

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

2.4.2 計算圖轉(zhuǎn)換到 ONNX

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

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

2.4.3 onnxruntime 和 caffe 的推理能力

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

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

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

2.5 端到端的模型轉(zhuǎn)換

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

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

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

3 總結(jié)

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

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