AI 框架部署方案之模型轉換

作者 | OpenMMLab（已授權） 

來(lái)源 | https://zhuanlan.zhihu.com/p/396781295 

編輯 | 極市平臺

導讀

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

1 模型轉換的意義

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

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

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

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

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

2 模型轉換的技術(shù)細節

2.1 計算圖生成

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

代碼語(yǔ)義分析：通過(guò)分析用戶(hù)代碼來(lái)解析模型結構，建立模型靜態(tài)表達。

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

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

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

2.2 計算圖中的自定義算子

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

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

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

2.3 目標格式（caffe/ONNX）

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

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

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

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

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

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

2.4 計算圖轉換到目標格式

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

2.4.1 計算圖轉換到 caffe

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

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

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

2.4.2 計算圖轉換到 ONNX

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

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

2.4.3 onnxruntime 和 caffe 的推理能力

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

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

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

2.5 端到端的模型轉換

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

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

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

3 總結

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

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