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

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

編輯 | 極市平臺(tái)

導(dǎo)讀

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

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

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

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

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

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

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

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

2.1 計(jì)算圖生成

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

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

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

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

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

2.2 計(jì)算圖中的自定義算子

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

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

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

2.3 目標(biāo)格式（caffe/ONNX）

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

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

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

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

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

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

2.4 計(jì)算圖轉(zhuǎn)換到目標(biāo)格式

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

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

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

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

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

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

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

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

2.4.3 onnxruntime 和 caffe 的推理能力

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

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

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

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

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

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

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

3 總結(jié)

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

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