PNNX: PyTorch 神經(jīng)網(wǎng)絡(luò)交換格式！
來源 | 知乎/?https://zhuanlan.zhihu.com/p/427620428
僅作為學(xué)術(shù)分享，如有侵權(quán)請聯(lián)系后臺刪除
作者 |?nihui@知乎
編輯?|?CV技術(shù)指南
【導(dǎo)讀】PNNX，全稱 pytorch neural network exchange，是一種開放的 pytorch 神經(jīng)網(wǎng)絡(luò)交換格式。PNNX 可以將 pytorch 訓(xùn)練的模型，導(dǎo)出為 PNNX 文件格式，更方便將模型部署到各個平臺和推理框架上。
為什么使用 PNNX，而不是 ONNX？
onnx 定義了一套標(biāo)準(zhǔn)化的 op，而各類深度學(xué)習(xí)算法迅猛發(fā)展，時常有 pytorch 的算子無法導(dǎo)出為 onnx，或?qū)С鰰r被 onnx 增加一大堆膠水 op，這些缺陷嚴(yán)重影響著 pytorch 模型的部署效率。究其本質(zhì)，是因為 onnx 致力于定義一套完全中立和理想化的算子表示，而pytorch 本身有其自身的 op 設(shè)計約定，兩者之間有著鴻溝。pnnx 則僅服務(wù)于 pytorch，可更好的表示 pytorch 的模型結(jié)構(gòu)，維持高層 op 對后端優(yōu)化也更加友好。
當(dāng)今流行的深度學(xué)習(xí)框架，實際只剩下 pytorch 和 tensorflow 兩家獨(dú)大，并不存在所謂 M種框架對 N種推理后端的情況，只有2種。目前，tensorflow 已布局 MLIR，作為其開放的模型交換格式；pytorch 雖然有 torchscript，但 torchscript 畢竟是一種類似 python 的語言，不適合作為一種模型交換格式，pnnx 的目標(biāo)就是補(bǔ)齊這一缺失，帶來比onnx 更加 pytorch 友好的方案。
PNNX
ok，現(xiàn)在就開始講 PyTorch 模型部署，我們現(xiàn)在這一頁看到的是全局的綜覽，是目前 PyTorch 模型部署的一些流程
我們這個 PPT 畫了橘紅色的虛線代表 PyTorch 的生態(tài)圈。然后 ONNX 用灰色的虛線畫出來是代表 ONNX 的生態(tài)圈。
以下是以我們一些常用的第三方的推理加速庫。
這是目前 PyTorch 部署的整體流程，基本上涵蓋了已知所有的方式，主要是 TorchScript 和 ONNX 兩種。
接下來我們會詳細(xì)的介紹各自這些部署方式的優(yōu)劣勢。
首先是我們看這個 PyTorch 到 libtorch 自有生態(tài)圈的部署
這種部署方式可以支持 PyTorch 模型里所有的 op，因為 libtorch 就是 PyTorch 的底層實現(xiàn)，所以只要 PyTorch 能跑的，導(dǎo)出來就一定能跑，然后 libtorch 也支持在 CPU 還有 GPU 上加速的，這是他的一個主要優(yōu)勢
但是他也有缺點(diǎn)，有兩個缺點(diǎn)。一個是他的速度沒有那些廠商專門優(yōu)化的庫快，比如，CPU 上 libtorch 的速度比 OpenVINO 慢，然后 NVIDIA GPU 上也比 TensorRT 會慢一點(diǎn)。另外這個 libtorch 的庫大小非常大，所以如果在移動端上部署的話，這個二進(jìn)制文件過于龐大可能有十幾兆這樣，所以也不太適合應(yīng)用于 APP
這是通過 ONNXRuntime 部署的方式
是由 PyTorch 先導(dǎo)出 ONNX 文件，然后使用 ONNX 自家的這個 ONNXRuntime 實現(xiàn)模型的推理
他這個方式一個好處是 ONNX 是支持其他一些訓(xùn)練框架，比如 TensorFlow、MXNet、Caffe，這樣你可以用同一個 ONNXRuntime 來部署任意一種訓(xùn)練框架導(dǎo)出的模型，這樣他的 workflow 是一致的
然后 ONNXRuntime 后端也有一些廠商優(yōu)化庫，比如 GPU 上可能會有 TensorRT，或者 OpenVINO 這樣優(yōu)化的后端，這樣他的執(zhí)行效率也會比較高
但是他有些缺點(diǎn)，比如 PyTorch 一些模型的算子可能在 ONNX 是沒有的，這樣導(dǎo)出 ONNX 的時候會出錯導(dǎo)不出來，以及 PyTorch 導(dǎo) ONNX，有時候會導(dǎo)出一個非常復(fù)雜的計算圖，這個情況會導(dǎo)致推理效率的下降
這里說的是第三方庫的部署方式
我們比較常規(guī)的方式是 PyTorch 導(dǎo) ONNX，然后再通過 ONNX 轉(zhuǎn)換到我們第三方推理框架的模型，比如我們 ncnn 也是通過 ONNX 來支持 PyTorch 的部署，那其他的比如 TensorRT、OpenVINO，還有騰訊的 TNN 也是通過 ONNX 來做這個部署
torchscript 那一條路呢，目前也看到些開源的框架支持通過 torchscript，比如 TRTorch、Forward，但是這條路目前用的比較少因為大家，ONNX 畢竟是時間比較久，用戶也比較多，資料也更多一點(diǎn)
使用第三方庫來做 PyTorch 部署優(yōu)勢，一個最主要的優(yōu)勢就是可以在你的目標(biāo)平臺上獲得最快的推理速度，TensorRT 在 NVIDIA 的 GPU 上最快，OpenVINO 在 Intel 的 CPU 上最快，那比如可能 ncnn 或者 TNN 在手機(jī)端 CPU 會更快一點(diǎn)
但是第三方庫有個更嚴(yán)重的問題就是，比如我剛才說的 ONNX，我也說了有一些 PyTorch 的算子導(dǎo)出 ONNX 是不支持的，而第三方庫相對 ONNX 可能支持的算子更加有限，有可能 ONNX 轉(zhuǎn)到第三方庫還有一部分算子是不支持，所以就只能支持一些比較常規(guī)的模型，復(fù)雜的話可能會導(dǎo)出不了不支持的情況
這就是目前 PyTorch 模型部署的三種主流的途徑吧
接下來我會介紹一個概念叫 lower
lower 這個概念其實是一個編譯器里面的概念，深度學(xué)習(xí)模型的轉(zhuǎn)換里面也是一個 lower 的概念
這 lower 什么意思啊就是一句話來說：就是用一個或多個框架支持的算子來模擬出模型里面所不支持的算子
下面就有兩個例子
比如我的模型里面寫了一個 padding=1 的卷積，那可能我的推理庫的卷積不支持有 padding
但是他支持一個單獨(dú)的叫 pad2d 的算子，這時候模型轉(zhuǎn)換過程中就會把 padding=1 的卷積
拆成兩個步驟，先做 padding，然后再做一個沒有 padding 的卷積，這個過程就是相當(dāng)于有兩個算子來實現(xiàn)了一個不支持的算子參數(shù)
下一個例子是 sigmoid
sigmoid 的操作，大部分機(jī)器或 GPU 上沒有獨(dú)立的函數(shù)叫 sigmoid，一般會拆成，先取負(fù)，然后 exp，再加1，最后1除tmp3 的操作，這樣就是個數(shù)學(xué)公式展開，相當(dāng)于用4個能實現(xiàn)的算子來模擬出 sigmoid。有的時候呢，這個平臺可能 exp 也不支持，那可能還會把 exp 用級數(shù)展開，用乘法加法的實現(xiàn)
這個過程，總之用多個支持的來模擬不支持的過程就叫 lower，這個箭頭從左到右就叫 lower
這個目的也是比較實在的，這樣我的后邊的推理庫只需要支持比較少量的算子，就能實現(xiàn)出很多很多其他的算子，對吧
然后如果我們的模型里面，因為我們的這個深度學(xué)習(xí)模型，他這個算子啊一直在發(fā)展，一直在變化，比如現(xiàn)在可能會出一個新的算子叫 swish，那 swish 激活其實是 x 乘以 sigmoid，如果我們有這種 lower 的方式的話，我們就可以把 swish 拆成 x 乘 sigmoid 的這樣兩步，這樣就可以覆蓋到更多新的算子，將來新的算子都可以兼容
第三個就是我們可以用這些支持的少量算子來兼容更多的訓(xùn)練框架，因為訓(xùn)練框架每個高層的算子可能在實現(xiàn)細(xì)節(jié)上有些不同，那我們拆開之后就可以去兼容這些不同的細(xì)節(jié)
這就是 lower 的一個概念
我們剛才 PyTorch 模型部署的流程中有個導(dǎo)出 torchscript 的過程，torchscript 的導(dǎo)出過程中也發(fā)生了這個 lower 的事情
比如，這里有個例子 F.normalize，這是一個比較常見的做一個 norm 的算子，這個算子在導(dǎo)出 torchscript 的時候，會因為 aten 沒有一個叫 normalize 的算子，所以會用 4 個 aten 的算子去模擬這個 normalize python 的一行
然后我們導(dǎo)出 ONNX 的時候，這里要明白 ONNX 并不是從 python 代碼直接轉(zhuǎn)到 ONNX，ONNX 導(dǎo)出時候是先通過導(dǎo)出 torchscript，然后從 torchscript ir 再轉(zhuǎn) ONNX
所以像左邊這個例子也是一個比較常見的操作叫 GroupNorm，這個 op 在 aten 和 torchscript 里面是存在的，所以導(dǎo)出到 torchscript 的時候，沒有發(fā)生 lower。但是 ONNX 沒有一個叫 GroupNorm 算子，所以使用了 6 個 ONNX 支持的算子去模擬了這個 GroupNorm，于是這個 GroupNorm 在 ONNX 被 lower 成了 6 個算子
經(jīng)過這兩次的 lower，我們會產(chǎn)生一個問題，就可能我們 PyTorch 里面那很簡單的一行，導(dǎo)出 ONNX 之后，經(jīng)過了兩層的層層 lower，最后會出來一個非常非常龐大的計算圖
這就是一個比較典型的例子，那這個龐大的計算圖會有比較嚴(yán)重的問題
首先，這個圖的算子過于細(xì)碎，太細(xì)碎的算子非常不利于推理的優(yōu)化，因為每個操作的力度都太細(xì)了
第二點(diǎn)是當(dāng)我們導(dǎo)出這么一個復(fù)雜的 ONNX 圖的時候，因為 ONNX 不是叫 Open Neural Network Exchange 嘛，ONNX 其實也是個模型文件嘛，我把這個模型文件給，比如我的第二位同事，他也是個算法工程師，他看到這個 ONNX 這個圖之后，再也無法對應(yīng)回原先的原始模型了，就是說你看到這個 ONNX 圖后，你不知道原始的網(wǎng)絡(luò)結(jié)構(gòu)是怎樣的了，因為這已經(jīng)變成一大坨漿糊了
第三個就是這個大的復(fù)雜計算圖里面有非常多的膠水 op，比如 ONNX 我們所熟知的 Gather、Unsqueeze op，這種 op 比如在 ncnn 里面是沒有做支持的，所以會有這個比如你用第三方庫，這樣 ncnn 里面去轉(zhuǎn)換模型時候會有轉(zhuǎn)不了的情況
針對這種情況呢，我們 ONNX 社區(qū)大缺弦老師有做了?onnx-simplifier，這個工具可以將比較大的復(fù)雜計算圖做一個簡化，去除里面一部分的那些膠水的 op。然后 ncnn 這邊的 ONNX 轉(zhuǎn)換工具還會再做一次這種簡化，把這種這么多的 op 再捏回到最初的那一個 op
這個把一些細(xì)碎的 op 捏回到一個大 op 的過程就叫圖優(yōu)化，這也是很多推理框架本身也會做的一些事情。因為 ONNX 這個圖就是非常復(fù)雜，所以幾乎每個推理框架都會做一些圖優(yōu)化
這個圖優(yōu)化工作主要做法就是模板匹配 pattern matching 對吧，然后再把匹配到的子圖用對應(yīng)的一個大 op 替換掉，這種圖優(yōu)化目前來說，還可以工作的很好，因為這個 pattern 我們可以寫一份工具去實現(xiàn)自動化的
但是呢，雖然表面上看上去金玉在外，敗絮其中啊。雖然說圖優(yōu)化能做一些這樣的好事，但是實際上也非常麻煩
麻煩在哪呢？
首先 ONNX 這個文件是個 protobuf，他本身沒有提供任何那種進(jìn)行圖優(yōu)化一些基礎(chǔ)庫的工作，不像 MLIR，所以我們工程師在寫圖優(yōu)化的時候，要寫大量的 if 判斷，參數(shù)判斷來實現(xiàn)圖優(yōu)化
第二點(diǎn)也是比較重要的是，我們 PyTorch 或者 ONNX 每次升級版本，或者 ONNX opset 從 9 變成 11 變成 13，他生成的導(dǎo)出這個 ONNX 的計算圖都會發(fā)生改變，那一改變之后呢，我們原先寫的圖優(yōu)化的 pattern 就不 work，他就無法再匹配到你新的。所以讓我們下游的這些推理框架開發(fā)者非常麻煩，就每次我們遇到版本升級之后，必須再寫一份圖優(yōu)化的一個函數(shù)來覆蓋到這個新的版本。嗯這個是一個永無止境的坑，因為我們知道 PyTorch 和 ONNX 永遠(yuǎn)在升級版本
第三點(diǎn)呢就是，有時候一些高層的這個 op 里面一些參數(shù)的變化也會導(dǎo)致圖的變化。比如我們所熟知的 PyTorch nn.Conv2d 卷積層，它里面有個參數(shù)叫 padding 模式，通常我們都是用常量 padding，但是它也支持 replicate 還有 reflect padding 的模式。那如果用這種參數(shù)的話，同樣是 Conv2d 導(dǎo)出的就會不一樣，這樣也會增加圖優(yōu)化匹配的復(fù)雜性
所以其實這個我們可以明白，就是我們當(dāng)初算法工程師寫 python 代碼的時候，其實是寫的一個比較簡單，比較干凈的 python 代碼，就是因為我們要經(jīng)過 torchscript 和 ONNX 中間商，通過他們導(dǎo)出之后，這個圖才會變得這么復(fù)雜。
那我們?yōu)槭裁匆欢ㄒ盟麄儗Π?，我們?yōu)槭裁床恢苯釉谠嫉?python 圖上，直接導(dǎo)出一個比較好的干凈的 ir 呢？
這里我思考了一些關(guān)于如何做一個比較好的模型交換格式的想法，因為模型交換嘛對吧
我們交換的話是應(yīng)該以人為本的，就是說我們希望這個模型文件出來之后，給另一位研究員看到的就是我們的 high-level 的
我們希望這個模型文件里面的 op 足夠 high-level，以便后端的廠商或者推理庫框架開發(fā)者能更好的做更激進(jìn)的優(yōu)化
我們也希望這個模型格式本身是對人類可讀和可編輯都比較友好的，像 torchscript 和 ONNX 里邊的二進(jìn)制是人類不可讀的，所以也是想要做一個方便人類去讀和編輯的表達(dá)形式
最后呢我認(rèn)為就是，我們主題叫 PyTorch Neural Network Exchange，我們選擇 PyTorch 的原因呢，也是因為目前深度學(xué)習(xí)訓(xùn)練框架中，PyTorch 在行業(yè)內(nèi)使用是最為廣泛的，所以在設(shè)計模型交換格式也是考慮我們只關(guān)注 PyTorch
所以今年，在2021年Q3 這是個剛新出爐的東西，叫 PNNX
那 PNNX 位置在整個大的圖里面在什么位置，在 torchscript 的下面的位置。我們剛才也說為什么不從 python 代碼直接導(dǎo)出 PNNX，還是要經(jīng)過 torchscript，因為 python 代碼其實是一個編程語言，不是一個可以解析的 ir，所以如果用 python 直接起的話會比較復(fù)雜，所以還是利用 torchscript，因為 torchscript 的畢竟是 PyTorch 自家生態(tài)，他稍微跟原始的 python 代碼會比較接近。然后通過 torchscript 里面的 ir 的轉(zhuǎn)換，再把這些 op 來捏回到原始的最高層這個 python 的 op
那這個箭頭是什么意思
這個箭頭是 torchscript 導(dǎo) PNNX，這個回去的箭頭代表 pt 轉(zhuǎn)出來后也會生成對應(yīng)的 python 源代碼，這個就是相當(dāng)于轉(zhuǎn)回去了，可以轉(zhuǎn)換成原始的 python，這樣就方便就是我們算法研究員去搭模型，對吧
然后右邊說這個 PNNX 會直接會導(dǎo)出 ncnn 的模型，這就是一個 PyTorch 模型部署的新途徑，是基于 PNNX 來實現(xiàn)
這里要說的一點(diǎn)是 PNNX 他不是一個新的 NN 算子標(biāo)準(zhǔn)
因為我們知道深度學(xué)習(xí)訓(xùn)練框架每個都有一個自己的算子定義，然后我們的后端推理庫或者平臺廠商優(yōu)化的 runtime，他也有自家的一套算子定義
然后針對這種情況，之前有些嘗試，比如 ONNX，Khronos NNEF，MLIR TOSA，TVM Relay。他們都會嘗試用一個大一統(tǒng)比較覆蓋更加全面的算子定義，來兼容所有的這些差異性。但是最后發(fā)現(xiàn)其實他們只不過又發(fā)明了一套新的算子定義而已，他們并沒有做到一個 Universal 的通用的大一統(tǒng)，最后只不過又做了一套新的，這不僅是增加了學(xué)習(xí)負(fù)擔(dān)，也是進(jìn)一步導(dǎo)致了算子之間、標(biāo)準(zhǔn)之間的差異和碎片化嘛
PNNX 這邊沒有發(fā)明新的算子標(biāo)準(zhǔn)
他的算子的定義是直接與 python 的代碼算子接口保持一模一樣，就是說他是直接利用這個 PyTorch API 就是這個算子。
我們可以看到我們寫的一個 nn.Linear，對吧。他轉(zhuǎn)成 PNNX 之后，模型里面寫的這個 op 就叫 nn.Linear。然后里面的所有參數(shù)都跟原始的 python api 保持一模一樣的名字。這樣的好處是，當(dāng)我拿到右邊的 PNNX 模型之后，我還可以通過他來轉(zhuǎn)回到原始的 PyTorch 的 python 代碼，因為算子定義和參數(shù)表達(dá)形式一模一樣，所以都可以轉(zhuǎn)回去。那這樣就相當(dāng)于是一個循環(huán)永動機(jī)，甚至你可以把轉(zhuǎn)回去的這個 model 直接拿來用或者直接拿來訓(xùn)練，或者你可以把它再導(dǎo)出成 PNNX，這樣就是相當(dāng)于一個完全一致的對應(yīng)關(guān)系
PNNX 的這個模型文件也是征詢了一些社區(qū)開發(fā)者的意見
因為我們 ncnn 的一些開發(fā)者對 ncnn 模型寫法都比較歡迎，因為他們會覺得 ncnn 文本的形式非常容易修改，非常容易魔改，比如加個 op，改個參數(shù)。所以 PNNX 這邊也是沿襲了 ncnn 的模型格式。這個格式跟 ncnn 基本是一樣的，但是擴(kuò)展了一些功能，就比如 ncnn 的參數(shù)列表，他的 key 只能是 0123 這種數(shù)字，那 PNNX 這邊就擴(kuò)展成可以用一個 string 作為 key，那這個和 python API 就可以保持對應(yīng)的關(guān)系
之前也考慮過?MLIR dialect，因為 MLIR 可以兼容萬物 op，對吧。但是 MLIR 玩過之后呢，我會覺得他 api 目前還不夠穩(wěn)定，MLIR 在 LLVM 的項目里也沒有正式發(fā)布一個穩(wěn)定版本，以及 MLIR 這個庫是需要用戶自己去 git clone LLVM 項目，然后自己去編譯的。這個事情會擋住很大一部分開發(fā)者，編譯也是比較困難。那還有一點(diǎn)就是 MLIR 我們知道其實是屬于 TensorFlow 生態(tài)，因為他是 Google 主導(dǎo)的。然后 PyTorch 和 TensorFlow 是完全兩個生態(tài)。如果基于 MLIR 去做的話，相當(dāng)于是在兩個生態(tài)之間牽線，這個事情并不是那么的好玩，所以還是使用了一個更加樸素、更加簡單的模型的形式
這是一些 PNNX 里面的特性
這里講的是，比如我 PyTorch 里面寫的一個簡單的算術(shù)表達(dá)式，這個表達(dá)式轉(zhuǎn)成 PNNX 之后會保留表達(dá)式的整體，他們不會被拆成加減乘除很多小算子。這樣我們閱讀模型本身時候會更加的方便，因為一看就能看出這是一整個表達(dá)式算術(shù)。另外在一些比如 GPU 或者可編程的硬件上，這種逐像素的操作，多次逐像素的操作他是可以合并出來，這樣減少這個層存儲訪問，可能效率會更高一點(diǎn)
這是 PNNX 第二個功能
我們左邊寫了一個 YOLOv5 里面的 Focus 模塊，當(dāng)他直接導(dǎo)出成 PNNX 的時候，下面是直接導(dǎo)出的一個狀況，他會導(dǎo)出成 4 個 slice 和 1 次 cat，這個圖是一個完全正確的表示形式
但是當(dāng)我們想要實現(xiàn)一個更高效的 Focus 的時候，我們其實不希望他拆成 5 個算子，我們希望讓他就一個步驟里去實現(xiàn)出整個 Focus 的過程，那這樣相當(dāng)于用一個大 op 來替換掉原先的一些小的碎 op
PNNX 也支持這種自動根據(jù) Module 來直接合并出一個大算子，當(dāng)我寫增加這個 moduleop=Focus 參數(shù)之后，我 PNNX 導(dǎo)出之后就不會再展開 Focus 的 Module，會把這個 Focus 一整塊當(dāng)成一個大 op 導(dǎo)出來
這是另一個特性
就是我們有時候 PyTorch 的一些算法會使用一些自定義 op，自己寫一個 cpp，或自己寫一個 cuda 的代碼實現(xiàn)自定義的 op，但是這種自定義 op 目前 ONNX 那邊是導(dǎo)不出來的。你肯定導(dǎo)不了，因為 ONNX 不可能有你這個自定義 op 的東西的
PNNX 是允許你導(dǎo)出這個自定義 op 的，我們加一個 customop 指向編譯出來的 torch so 之后，他就會把這個自定義 op 導(dǎo)出成 PNNX.custom_op，然后自定義 op 的名字，然后自定義 op 的參數(shù)，比如這個 slope、scale，他都會寫在后面，就是參數(shù)也會寫在后面。
但是具體這個自定義 op 的實現(xiàn)他是沒有的，因為你的實現(xiàn)是 cpp 和 cuda 的代碼，對吧。但是 PNNX 只能負(fù)責(zé)這個計算圖的導(dǎo)出，后面的話，比如推理庫那邊可能會支持一些自定義 op 的插件，就可以直接自定義 op 實現(xiàn)了，這樣也相當(dāng)于是解決了自定義 op 導(dǎo)出的問題
這是另一個主要特性
這是 PyTorch QAT，就是量化感知訓(xùn)練的算子。目前來說，如果你用 ONNX 去導(dǎo)出的話，他可能只支持部分 QAT 的訓(xùn)練策略
然后 PNNX 這邊，如果你對量化模型導(dǎo)出，會嘗試做一個比較好、比較完整的量化相關(guān)的一些參數(shù)的處理，比如量化的一些 Quantize、Conv2d 這些層，他會把量化的 scale，還有 zeropoint 這些參數(shù)記錄下來，以及我們量化的 Conv 的 weight，它也會存成對應(yīng)的 int8 數(shù)據(jù)，存在 bin 里，還有 weight 的 per-channel 的 scale，zeropoint 也會記錄，這樣就解決了導(dǎo)出 QAT op 方法的一個問題，這個也是目前 ONNX 做的還不夠好的一方面
PNNX 在轉(zhuǎn) torchscript 導(dǎo)出的時候可以寫 shape
如果你寫下 inputshape 之后，導(dǎo)出的 PNNX 后面會產(chǎn)生每個 feature blob，對應(yīng)的 shape 信息會寫在后面，這樣會幫助到一些有 shape 相關(guān)信息參與的表達(dá)式的常量折疊優(yōu)化，有時候這種什么 reshape 呀這種的后面參數(shù)，就會通過我們 shape 的推斷和傳播把它變成一個常量。
有的時候我們的模型支持動態(tài) shape，比如這里的 YOLOv5，這種情況我們的工具允許你再寫第二個 shape，第二個 shape 可以跟第一個 shape 不一樣。那你不一樣的時候呢，導(dǎo)出的模型里面就會帶有問號，這個 shape 里帶有的問號，這位置就代表當(dāng)前的一維，比如 w h 兩個維度，他的尺寸是不固定。這個可能會影響一些后端的優(yōu)化的策略，但從 PNNX 設(shè)計上說，支持靜態(tài) shape，也支持動態(tài) shape
這里稍微提一下這個 PNNX 內(nèi)部的圖優(yōu)化的技術(shù)
一個庫，一個基礎(chǔ)設(shè)施，ONNX 是沒有的，對吧。PNNX 剛才也說是從 torchscript 再捏回到原先高層 op，所以還是需要做?torchscript ir?的圖優(yōu)化
這里 PNNX 寫了這個叫 GraphRewritter 的類，只要指定我們的 pattern 的 ir，指定好后，就會自動從 torchscript 里找到一個匹配的封閉子圖，然后把它替換成我們的 target 的 op，最后這個 target op 的一些參數(shù)也可以，從匹配的子圖里去獲取到原先那個參數(shù)，然后寫到 target 參數(shù)里，PNNX 代碼里面就有大量的捏回原始高層算子的這個過程，所以有一個專門做的基礎(chǔ)設(shè)施，方便做捏算子的圖優(yōu)化過程
這是 PNNX 代碼的整體架構(gòu)
首先是加載 torchscript，這是用 libtorch 的 api 完成的，里面會做一些 inline 或者 shape 的推斷傳播，這部分做完后還是 torchscript ir
然后從 torchscript ir 轉(zhuǎn)換成 PNNX 的 ir?數(shù)據(jù)結(jié)構(gòu)，然后在這個數(shù)據(jù)結(jié)構(gòu)上去做捏算子，就是把 torchscript ir 一個個捏成 nn 的 Module，或者是 torch 的 function，或是 Tensor 的 function。后面也會做一些比較常規(guī)的優(yōu)化，比如廢棄代碼的刪除，或者一些表達(dá)式的折疊這種
那當(dāng) PNNX IR 數(shù)據(jù)結(jié)構(gòu)優(yōu)化完成之后呢，可以直接產(chǎn)生 python，因為這個 ir 里面已經(jīng)是最高層的，所以 python 代碼我們在轉(zhuǎn)的時候，基本就是1比1的寫出來就可以
當(dāng)我要轉(zhuǎn)成 ncnn 的時候，還需要 ncnn 的對應(yīng)的算子轉(zhuǎn)換過程。這里需要去寫從 nn 的 Module 轉(zhuǎn)換成 ncnn 對應(yīng) op 的轉(zhuǎn)換器，每個算子也會實現(xiàn)一個，也會針對 ncnn 模型的特殊的特征加一些 Split 這種層
這就是 PNNX 整體框架的架構(gòu)設(shè)計
這個代碼在哪里呢
這個代碼現(xiàn)在目前在 ncnn 的一個 pr 里邊。如果聽眾朋友們想要交流的話，可以直接在 pr 里留下你們的意見
目前 PNNX 兼容 PyTorch 1.8、1.9 和 1.10 版本，也可以導(dǎo)出 159 種 PyTorch 上層 op，其中有 80 個可以導(dǎo)出成 ncnn 對應(yīng)的 op，也做了自動單元測試，還有代碼覆蓋率，目前是 79%，也還行，因為這個也是一個比較剛開發(fā)不久的項目該 79。目前常用的 CNN 模型，像 resnet，shufflenet 這種模型，都是可以完美工作，就是 torchscript，轉(zhuǎn) PNNX、轉(zhuǎn) python 或 ncnn 這些都可以正常搞定，然后出來的推理結(jié)果跟原始的 python 是一模一樣的
后面呢會打算支持更多的 PyTorch 算子，增加更多的單元測試，增加一些端到端 RNN 或者 Transformer 的模型測試，還會寫一些使用教程或者開發(fā)文檔。因為現(xiàn)在一直致力于 coding，所以這些還比較欠缺。
往期精彩：
?時隔一年！深度學(xué)習(xí)語義分割理論與代碼實踐指南.pdf第二版來了！