大家好。我叫Manav Dalal，今天我將講解如何通過(guò)新的ONNX導(dǎo)出器簡(jiǎn)化模型導(dǎo)出流程。如果你還沒(méi)有聽說(shuō)過(guò)ONNX，它是一種用于表示機(jī)器學(xué)習(xí)模型的開放格式。它定義了一套通用的運(yùn)算符，機(jī)器學(xué)習(xí)和深度學(xué)習(xí)模型的基本原理，并提供了一個(gè)通用的文件格式，因此你可以保存并與各種工具一起使用,這樣你就可以在你喜歡的框架中進(jìn)行開發(fā)而且不用太擔(dān)心與其相關(guān)的推理邏輯。

另一個(gè)很棒的地方是它真正解鎖了許多硬件優(yōu)化，無(wú)論你是在處理網(wǎng)絡(luò)、NPU、GPU、CPU，還是其他任何設(shè)備，你都可以從ONNX訪問(wèn)你的模型。這是一些可以導(dǎo)出onnx的合作伙伴：

簡(jiǎn)單回顧一下自PyTorch 1.2以來(lái)的情況。我們使用torch.onnx.export API導(dǎo)出到Onnx模型。這個(gè)舊api有很多問(wèn)題，通過(guò)Torch Script tracing(torch.jit.trace)可以在模型執(zhí)行時(shí)生成靜態(tài)圖。雖然這樣可以工作，但靜態(tài)圖存在很多限制。如果您有控制流程如if語(yǔ)句和循環(huán)，它在處理這些內(nèi)容時(shí)就會(huì)表現(xiàn)很差。此外，它也不能很好地處理 training 和 eval 之間的細(xì)微差別。它也無(wú)法處理真正動(dòng)態(tài)的輸入。

現(xiàn)在，我們已經(jīng)有了一些解決方案如torch.jit.script（比如對(duì)控制流問(wèn)題的解決）Torch.jit.script本身是Python的一個(gè)子集，因此雖然有一對(duì)一的映射，但并非所有內(nèi)容都有映射…因此，有時(shí)會(huì)忽略諸如原地操作之類的東西。

但接下來(lái)我們有了新的方式——Torch Dynamo。簡(jiǎn)而言之，在2.0版本中，它被視為一種JIT編譯器，旨在加速未修改的PyTorch程序，從而實(shí)現(xiàn)可用性和性能的提升。它通過(guò)重寫Python字節(jié)碼，在FX圖中獲取PyTorch操作的部分，并使用它來(lái)獲得我們想要的所有這些巨大的性能提升。它也是Torch Compile、Torch Export的基礎(chǔ)引擎，

今天我要談?wù)摰氖切碌腛NNX Exporter。我們推出的一個(gè)目前還在測(cè)試階段的新導(dǎo)出器。也許你們中的一些人已經(jīng)使用過(guò)它，也可能有一些人還沒(méi)用過(guò)。今天我在這里希望能夠說(shuō)服你們至少去試一試。它看起來(lái)有點(diǎn)類似于舊的API，但稍微復(fù)雜一些。你可以使用你的Torch.nn模塊以及與該模塊相關(guān)的參數(shù)進(jìn)行導(dǎo)出（無(wú)論是位置參數(shù)還是關(guān)鍵字參數(shù)）。還有一個(gè)新的參數(shù)導(dǎo)出選項(xiàng)，你可以直接添加一些相對(duì)具體的參數(shù)。通過(guò)利用Dynamo來(lái)捕捉graph，我們可以更容易地生成與該圖形等效的ONNX，并且可以保留這些模型的動(dòng)態(tài)特性。相對(duì)于以前的靜態(tài)圖，現(xiàn)在我們可以擁有動(dòng)態(tài)圖，我們可以用它們來(lái)獲得一系列的好處。此外，onnx函數(shù)現(xiàn)在被用于封裝ATen運(yùn)算符和nn.modules，這對(duì)于導(dǎo)出后的onnx轉(zhuǎn)換和lowering操作非常有用，可以去除復(fù)雜性，直接在onnx模型中保留更多語(yǔ)義和結(jié)構(gòu)信息，這可以在許多常見的優(yōu)化場(chǎng)景中用于子圖模式匹配。讓我給你展示一個(gè)的例子。

你可以在這里看到不同的圖表。左側(cè)是舊的靜態(tài)圖表，這是可以的，在導(dǎo)出模型時(shí)工作正常，并且在某種程度上也經(jīng)過(guò)了優(yōu)化。但是，在中間，你可以看到現(xiàn)在的torch dynamo圖表是什么樣的。有更多的信息，并且所有的操作符都被捕捉在函數(shù)中。這真的很有用，因?yàn)閷?shí)際上它們內(nèi)聯(lián)之后，你可以看到捕捉到了更多的信息。支持了許多數(shù)據(jù)類型。這是所有操作符支持的數(shù)據(jù)類型，而不僅僅是你用于導(dǎo)出模型的那個(gè)數(shù)據(jù)類型。這對(duì)于之后的優(yōu)化非常有用。你只會(huì)得到更多的信息。

新的torch.onnx.dynamo導(dǎo)出API通過(guò)可選的ExportOptions支持配置。值得注意的屬性包括：

.onnxregistry：配置ATen分解，指定ATen IR到onnx運(yùn)算符的映射，并提供API來(lái)支持三種主要的自定義運(yùn)算場(chǎng)景。
-處理不支持的ATen運(yùn)算符（缺少的運(yùn)算符）。
- custom operators 但是有已存在的onnx runtime support。
- custom operators 但是沒(méi)有已存在的onnx runtime support。

.fake context：torch.onnx.enablefakemode()的結(jié)果是一個(gè)上下文管理器，它啟用了導(dǎo)出大規(guī)模模型的虛擬模式輸入支持，這些模型通常會(huì)在使用TorchScript導(dǎo)出時(shí)內(nèi)存溢出。

diagnostic options：導(dǎo)出器中深度融入了一個(gè)新的診斷系統(tǒng)，以提高調(diào)試能力和自助服務(wù)能力。

它返回一個(gè)ExportOutput對(duì)象，通過(guò)該對(duì)象可以進(jìn)一步通過(guò).model proto屬性進(jìn)行推理或轉(zhuǎn)換，或者通過(guò).save(path)方法直接將內(nèi)存中的onnx模型序列化到磁盤以方便使用。

因此，讓我們談?wù)勎覀冮_發(fā)的這個(gè)新API的亮點(diǎn)，希望你們能夠使用。正如我之前提到的，我們有可選的導(dǎo)出選項(xiàng)，并且你可以用它做很多事情。首先，有onnx注冊(cè)表，它允許你配置Aten分解，你可以為Onnx運(yùn)算符預(yù)設(shè)AtenIR并進(jìn)行映射，并提供API用于處理一些常見情況，比如處理不支持的A10運(yùn)算符、具有現(xiàn)有OR支持的自定義運(yùn)算符以及在Onnx運(yùn)行時(shí)不支持的自定義運(yùn)算符。因此，在前兩種情況下，你可以提供這些運(yùn)算符，并在Onnx腳本中自己編寫它們（稍后會(huì)介紹）。這有點(diǎn)像一門語(yǔ)言，但是你基本上可以編寫這些語(yǔ)言或運(yùn)算符，并指定它們?nèi)绻豢捎脮r(shí)應(yīng)該如何編寫出來(lái)。在第三種情況下，你不僅可以編寫自定義運(yùn)算符，還可以編寫所需注冊(cè)的ORT內(nèi)核。

我們還有fake上下文，這要?dú)w功于meta的虛假tensor。我們可以在不加載所有權(quán)重和運(yùn)行模型的情況下導(dǎo)出模型。你知道，計(jì)算在如今是十分昂貴的，而且很難總是獲得足夠的計(jì)算資源來(lái)完成我們需要使用這些模型的任務(wù)，但是有了fake上下文，我們就能夠做到。我們可以在所需設(shè)備準(zhǔn)備好可以使用的ONNX模型，而無(wú)需花費(fèi)可能多達(dá)幾分鐘的時(shí)間導(dǎo)出該模型。我們知道之前在大型模型方面，特別是最近大家都在談?wù)摰腖LAMA模型方面，使用Torch腳本導(dǎo)出器時(shí)會(huì)遇到問(wèn)題。這是一個(gè)龐大的模型，是的，我們會(huì)遇到內(nèi)存問(wèn)題，但是有了fake上下文，你就不必遇到這些問(wèn)題了。我們還注意到，以前的Torch腳本導(dǎo)出器在調(diào)試能力和自助能力方面并不好，因此，在現(xiàn)在有很多診斷和自助能力選項(xiàng)，所以當(dāng)你遇到錯(cuò)誤時(shí)（因?yàn)楝F(xiàn)在仍處于測(cè)試階段），很容易弄清楚該怎么做或如何解決這些錯(cuò)誤。這是我們非常關(guān)注的問(wèn)題，我們確實(shí)在意這個(gè)，現(xiàn)在，在你選擇的模型上運(yùn)行導(dǎo)出之后，你會(huì)得到ExportOutput對(duì)象。這是一個(gè)存在于內(nèi)存中的對(duì)象，你可以進(jìn)行推理并通過(guò)獲取的模型 proto 屬性來(lái)進(jìn)行操作。因此，你可以查看它，會(huì)有很好的選項(xiàng)來(lái)對(duì)他進(jìn)行更改和性能更新。當(dāng)然，你也可以像保存其他 Onnx 模型一樣將它保存到磁盤上，這個(gè)過(guò)程也非常簡(jiǎn)單?？傊?，效果非常好。OnnxScript 是我之前提到的用于實(shí)現(xiàn)新操作的方式。幾個(gè)月前，它作為一個(gè)開放的倉(cāng)庫(kù)宣布出來(lái)。你可能已經(jīng)讀過(guò)相關(guān)博客，也可能沒(méi)有。簡(jiǎn)而言之，它是一種符合慣用方式的簡(jiǎn)單方法來(lái)編寫ONNX函數(shù)。

根據(jù)我的經(jīng)驗(yàn)，在以前使用Onnx時(shí)，這實(shí)際上是一件很具挑戰(zhàn)性的事情。我知道，像編寫自己的運(yùn)算符一樣，這是一種痛苦的過(guò)程。但現(xiàn)在它非常簡(jiǎn)單明了。這是一個(gè)希望你能夠閱讀的函數(shù)操作符。如果你可以，你會(huì)發(fā)現(xiàn)代碼相當(dāng)直觀。這要?dú)w功于OnnxScript讓它變得非常容易。它非常簡(jiǎn)單明了，易于操作，而且對(duì)于那些不太精通技術(shù)的人來(lái)說(shuō)也很易懂。所以整個(gè)過(guò)程都很容易上手。如果你選擇，你也可以直接使用OnnxScript完全編寫Onnx模型。這取決于你。但是，就像今天的演講中的其他人提到的一樣，我們將直接與Torch.export進(jìn)行集成，

另外一個(gè)我提到的事情是，Onnx轉(zhuǎn)換和優(yōu)化工具將很快推出，以幫助利用這些模型。我們希望這個(gè)工具能被很多運(yùn)行環(huán)境使用，謝謝。下列是一些參考資料：

Installrequiredpackages:

pipinstalltorchonnxscriptonnxruntime
ExportaPyTorchModeltoONNX:

Documentation:aka.ms/pytorchtoonnx
ExtendingtheONNXRegistry:

Documentation:aka.ms/onnx/registry
APIDocsforONNXDynamo:

Documentation:aka.ms/onnx/dynamo
IntroducingONNXScript:

Blog:aka.ms/onnxscript/blog1
ONNX:

Officialwebsite:aka.ms/onnx
ONNXRuntime:

Documentation:aka.ms/onnxruntime
ONNXScript:

Documentation:aka.ms/onnxscript

審核編輯：湯梓紅