【GiantPandaCV導語】torch.fx對于PyTorch來說確實是一個比較好的工作，因為它消除了一些動態(tài)圖和靜態(tài)圖的Gap。比如在圖改寫方面，torch.fx讓PyTorch想做一些其它靜態(tài)圖框架的算子融合優(yōu)化非常容易。并且torch.fx讓后訓練量化和感知訓練量化以及AMP等的實現(xiàn)難度大大降低，這得益于我們可以直接在Python層操作這個IR，所以我認為這是一個不錯的工作。尤其是對使用PyTorch開發(fā)的算法工程師來說，現(xiàn)在可以基于這個特性大開腦洞了。我之前圍繞FX也做了一個QAT的工作，感興趣可以閱讀：基于OneFlow實現(xiàn)量化感知訓練。torch.fx的賣點就是，它使用純Python語言實現(xiàn)了一個可以捕獲PyTorch程序的計算圖并轉化為一個IR的庫，并且非常方便的在這個IR上做Pass，同時提供將變換后的IR Codegen合法的Python代碼功能。我覺得算是達到了在Eager下寫Pass就像做鏈表插入刪除題目一樣順滑。

0x0. 動機

最近沐神在B站上分享了一些經(jīng)典論文的閱讀視頻，我也跟著看了幾個比如TransFormer，VIT等，很贊。所以我就想試試沐神這種論文閱讀方法，找一篇論文來閱讀一下。目前我比較關注的是工程方向的論文，正好上周PyTorch 放出了torch.fx的論文，所以我就以這篇論文為例來試試。沐神的論文閱讀方法大概是這樣（數(shù)字代表先后順序）：

1. 標題
2. 摘要
3. 引言
4. 結論
5. 相關工作
6. FX特性
7. 實驗
8. 評論

PyTorch FX論文的鏈接在：https://arxiv.org/pdf/2112.08429.pdf 。下面我就以沐神的論文閱讀順序來分享一下閱讀體驗，幫助大家搞清楚PyTorch FX這個特性到底是什么，以及它可以在PyTorch中發(fā)揮什么作用。

0x1. 標題

我們可以這樣進行翻譯：TORCH.FX：基于Python的深度學習實用程序「捕獲」和「變換」。這個基于Python的深度學習實用程序這里就可以理解為基于PyTorch開發(fā)的模型程序了，然后這里的重點是捕獲和變換?，F(xiàn)在還不清楚這里的捕獲和變換是什么意思，我們可以接著往下讀。插一點言外話，我半年前就有關注FX，之前在OneFlow框架下也成功接入了FX并且做了一個QAT的任務，而那個時候FX還是沒有論文的，所以我感覺這篇論文更像是一個對TORCH.FX這個特性的總結以及占位的那種感覺。

0x2. 摘要

摘要部分簡單指明了像PyTorch這種基于動態(tài)圖執(zhí)行模式的深度學習框架雖然提升了用戶的易用性。但是在一些真實場景中，用戶可能是需要捕獲和變換程序結構（也可以直接理解為神經(jīng)網(wǎng)絡的結構）來進行性能優(yōu)化，可視化，分析和硬件調優(yōu)等。為了解決這個痛點，PyTorch設計了torch.fx這個模塊來做PyTorch程序的捕獲和變換，這個模塊是純Python開發(fā)的。

這一節(jié)主要是講了一下torch.fx的賣點，就是說動態(tài)圖雖然易用性很強，但是圖結構不能被提前感知和變換，但通過這篇論文的torch.fx模塊，這件事就成啦！

0x3. 引言

早期的圖模式或者叫define-and-run的靜態(tài)圖框架有Caffe，TensorFlow等，它們設計了一個表示圖的IR，用戶通過調用這些框架提供的API來構建IR。然后我們可以在這個IR上做程序微分，將IR切分到設備上實現(xiàn)并行，量化，性能優(yōu)化等等。但這些事情一般都要求開發(fā)者在領域特定的語言上去做，比如以OneFlow的靜態(tài)圖模式為例，要做圖切分，量化，性能優(yōu)化等都需要基于C++去開發(fā)，并且調試也會相對比較難（要借助pdb，gdb等等專業(yè)工具）。

現(xiàn)在的eager模式或者叫define-by-run的動態(tài)圖框架有PyTorch，TensorFlow Eager模式等，它們可以隨心所欲的讓用戶基于腳本語言編程并且可以解決大多數(shù)的訓練（基于自動微分）和預測任務。但是有一些變換比如「量化和算子融合」 是不能直接做的，而這一點在靜態(tài)圖模式下則很簡單。為了消除這種Gap，動態(tài)圖框架需要一種從用戶的程序捕獲圖結構的方法來使能這些變換。

實際上，這種程序捕獲技術在PyTorch中早就有了，那就是TorchScript，它基于Python程序的AST來構造IR，對整個Python程序進行全面建模。但這樣做仍存在一個問題，那就是程序捕獲的技術復雜性太大，并且在這個高復雜性的IR上寫變換太難了。作為對比，我們可以將需求簡化一下，從全面建模Python程序變成只需要建模到可以做量化和算子融合等變換即可。要做到這兩件事，我們其實只需要程序中的那個DAG結構就可以了，而不需要程序中隱藏的更高層的API的結構（比如卷積和BN），這句話的意思就是說我們可以不關注卷積和BN這種可能由多種操作拼出來的nn.Module，我們在他的上層進行截斷就行，即最后獲得了一個high-level API的DAG。

基于上面的思路，torch.fx被提出，它關注深度學習程序中的DAG并且提供了定制接口來獲取這個DAG。這樣torch.fx就可以實現(xiàn)大多數(shù)深度學習框架中的圖變換，同時還提供了一套簡單易用的APIs來幫助用戶自定義變換?？偨Y一下，torch.fx的核心賣點是：

1. 對深度學習程序很重要的程序捕獲和轉換的實用分析特性。Trace
2. 一個僅用Python實現(xiàn)的程序捕獲庫，可以進行定制以捕獲不同級別的程序細節(jié)。Pure Python
3. 一個簡單的只有 6 條指令的 IR，用于表示捕獲的程序，重點是易于理解和易于進行靜態(tài)分析。IR
4. 用于將變換后的代碼返回到宿主語言生態(tài)系統(tǒng)的代碼生成系統(tǒng)。Codegen
5. 關于如何在實踐中使用 torch.fx 開發(fā)性能優(yōu)化、程序分析、設備lowering等功能的案例研究。Eager Pass

0x4. 結論

我們提出了torch.fx，這是一個用于捕獲和轉換PyTorch程序的純Python系統(tǒng)。我們分析了使相關系統(tǒng)復雜化的因素，包含控制流，可變性和數(shù)據(jù)模型，并展示了 torch.fx 如何通過關注常見用例和可定制性來避免復雜性。我們在優(yōu)化、分析和設備lowering方面調查了 torch.fx 的各種用例，并展示了 torch.fx 的 API 設計如何實現(xiàn)這些結果。

0x5. 相關工作

在捕獲和變換程序時，eager和graph模式的深度學習框架都必須在「捕獲程序結構」、「程序特化」和「保存程序的IR的設計」方面做出選擇。這些選擇的組合決定了可在框架中表示的「程序空間」、「編寫變換的難易程度」以及「生成的變換程序的性能」。「一般來說，支持程序的高性能運行需要更復雜的捕獲框架和IR，從而使轉換更難編寫」。每一段相關工作我就不詳細過了，只描述每一段工作的核心是在說什么，相關細節(jié)大家可以查看原始論文。

0x5.1 捕獲程序結構

這一節(jié)提到了PyTorch的jit.trace，MxNet Gluon，TensorFlow的tf.function等程序捕獲方法，并指出這些方法只能處理Python的一些子集。然后，TorchScript通過在AST上分析可以處理控制流和更多的Python語法。然后還提了一下Julia和Swift For TensorFlow中將捕獲程序結構的接口集成到了非Python的宿主語言中，要使用的話需要用戶放棄Python生態(tài)系統(tǒng)。

0x5.2 特化程序

對于a+b這個Python語句來說，這個表達式對a和b的類型沒有限制。但當深度學習框架捕獲程序時一般會對這兩個變量進行特化，以便它們只對特定類型或者張量有效。在深度學習框架中處理的大多數(shù)程序都是特化類型的程序，特化程度越高，能夠處理的輸入就越少。例如torch.jit.trace在執(zhí)行trace的時候只能處理某些擁有合法輸入shape的輸入。接下來還討論了LazyTensor和Jax的jit來說明為了更好的處理特化程序中捕獲的失敗，它們做了哪些努力。

0x5.3 IR設計

深度學習框架都有自己的IR設計，Caffe和TensorFlow使用Protocol Buffers格式。而PyTorch和MxNet使用C++數(shù)據(jù)結構來表示IR并額外綁定到Python。這些IR設計在runtime階段表現(xiàn)都會比較好并且可以統(tǒng)一被序列化。但從另外一個角度來說，這些IR表示相比于純Python語言的表示都需要更高的學習成本。接下來，這一節(jié)討論了控制流和狀態(tài)的問題，用來表明要處理這些問題需要設計較為復雜的IR以及要基于這個IR做較為復雜的分析才行。

基于上面幾點，論文提出了torch.fx的基本設計原則：

• 避免支持長尾分布，復雜的樣例。主要關注經(jīng)典模型的程序捕獲和變換。
• 使用機器學習從業(yè)者已經(jīng)熟悉的工具和概念，例如Python的數(shù)據(jù)結構和 PyTorch 中公開記錄的算子 。
• 使程序捕獲過程具有高度可配置性，以便用戶可以為長尾需求實現(xiàn)自己的解決方案。

這一節(jié)主要對一些相關工作進行了展開，以此來突出torch.fx的核心賣點，就是說我雖然不能像TorchScript這樣的IR處理一些比較難的Case（比如動態(tài)控制流），但是我在神經(jīng)網(wǎng)絡這個領域里做得夠用就可以了。最關鍵的是我的實現(xiàn)很簡單，是純Python的庫，這樣用戶寫變換就會很簡單，學習成本會很小并且易用。（簡單不代表不強大！

0x6. FX特性

以簡單為基本原則，torch.fx通過符號跟蹤來捕獲程序，并通過一個簡單的6個指令的IR來表示它們，并基于這個IR重新生成Python代碼來運行它。為了避免JIT特化中的重捕獲的復雜性，torch.fx沒有對程序本身進行特化，而是依靠變換來決定在捕獲期間需要實現(xiàn)哪些特化。用戶也可以配置符號跟蹤的過程來實現(xiàn)自定義捕獲需求。

Figure1給我們展示了使用torch.fx.symbolic_trace捕獲程序的例子，輸入可以是一個torch.nn.Module或者函數(shù)，并且捕獲之后的結構被存在一個Graph對象里面。該Graph對象和GraphModule中的模塊參數(shù)相結合，GraphModule是 torch.nn.Module 的子類，其 forward 方法運行捕獲的 Graph。我們可以打印此圖的Nodes以查看捕獲的 IR。placeholder節(jié)點表示輸入，單個output節(jié)點表示Graph的結果。call_function 節(jié)點直接引用了它將調用的 Python 函數(shù)。call_method 節(jié)點直接調用其第一個參數(shù)的方法。Graph 被重組為 Python 代碼（traced.code）以供調用。

Figure2展示了使用 torch.fx 進行變換的示例。變換是找到一個激活的所有實例并將它們替換為另一個。在這里，我們使用它來將gelu 替換 relu。

0x6.1 程序捕獲

torch.fx的符號跟蹤機制使用一個Proxy數(shù)據(jù)結構來記錄給定一個輸入之后經(jīng)過了哪些Op。Proxy是一個duck-typed類型的Python類記錄了在它之上的的屬性訪問和調用方法，是程序中真實Op的上層抽象。duck-typed可以看一下這里的介紹：https://zh.wikipedia.org/wiki/%E9%B8%AD%E5%AD%90%E7%B1%BB%E5%9E%8B 。PyTorch的算子以及Python子集的某些函數(shù)都會被這個Proxy包裝一次，然后在符號跟蹤傳入的是一個nn.Module時，會對這個nn.Module中的子nn.Module也進行Proxy包裝，當然還包含輸入數(shù)據(jù)。這樣程序中的輸入和其它Op都是duck-typed類型的Proxy對象，我們就可以執(zhí)行這個程序了，也就是符號跟蹤的過程。符號跟蹤的過程通過一個Tracer類進行配置，它的方法可以被重寫以控制哪些值被作為Proxy對象保留，哪些值被unpack。（Proxy記錄下來的Op可以進行unpack，unpack之后可以拿到真實的Tensor, Parameter和運算符等等）。通過Proxy和Tracer類的配合，torch.fx就可以完成PyTorch程序的符號跟蹤，需要注意的是這里的符號跟蹤的意思就是運行一遍這個被代理之后的nn.Module的forward。

0x6.2 中間表示

torch.fx的中間表示（IR）由一個Python數(shù)據(jù)結構Graph來做的。這個Graph實際上是一個包含一系列Node的線性表。節(jié)點有一個字符串操作碼opcode，描述節(jié)點代表什么類型的操作（操作碼的語義可以在附錄 A.1 中找到）。節(jié)點有一個關聯(lián)的目標，它是調用節(jié)點（call_module、call_function 和 call_method）的調用目標。最后，節(jié)點有 args 和 kwargs，在trace期間它們一起表示 Python 調用約定中的目標參數(shù)（每個opcode對應的 args 和 kwargs 的語義可以在附錄 A.2 中找到）。節(jié)點之間的數(shù)據(jù)依賴關系表示為 args 和 kwargs 中對其他節(jié)點的引用。

torch.fx 將程序的狀態(tài)存儲在 GraphModule 類中。GraphModule 是轉換程序的容器，暴露轉換后生成的代碼，并提供 nn.Module 類似的參數(shù)管理APIs。GraphModule 可以在任何可以使用普通的 nn.Module 的地方使用，以提供轉換后的代碼和PyTorch 生態(tài)系統(tǒng)的其余部分之間的互操作性。

0x6.3 源對源的變換

torch.fx 變換pipline的最后階段是代碼生成。torch.fx 并不是退出 Python 生態(tài)系統(tǒng)并進入定制的運行時，而是從變換后的 IR 生成有效的 Python 源代碼。然后將此變換后的代碼加載到 Python 中，生成一個可調用的 Python 對象，并作為forward方法安裝在 GraphModule 實例上。使用代碼生成允許將 torch.fx 變換的結果安裝在模型中并用于進一步的變換。例如，在圖3中，我們拿到trace原始程序的結果并將其安裝為新模塊中的激活函數(shù)。

到這里PyTorch FX特性就精讀完了，但查看FX的論文可以發(fā)現(xiàn)還有一節(jié)叫作Design Decisions，分別介紹了Symbolic Tracing，Configurable Program Capture，AoT Capture without Specialization，Python-based IR and Transforms等等FX實現(xiàn)中依賴的一些想法和 決策，以及它們的好處等。我理解這一節(jié)就是Introduction的加強版，所以就不繼續(xù)講解這一小節(jié)了，如果你擔心會錯過什么細節(jié)知識可以閱讀論文原文。

0x7. 實驗

torch.fx的一個目標就是簡化深度學習模型產生的IR，下面的Figure5以ResNet50為例展示了TorchScript IR和torch.fx IR的差別，相比于TorchScript IR，torch.fx IR確實簡單并且可讀性更強。

我們知道后量化以及量化感知訓練可以提示程序推理時的性能，下面的Figure6就展示了基于torch.fx實現(xiàn)的后量化（使用FBGEMM量化算子）應用在DeepRecommender模型之后，在Intel Xeon Gold 6138 CPU @2.00GHz上的性能表現(xiàn)?；?code style="font-size: 14px;word-wrap: break-word;margin: 0 2px;background-color: rgba(27,31,35,.05);font-family: Operator Mono, Consolas, Monaco, Menlo, monospace;word-break: break-all;color: #3594F7;background: RGBA(59, 170, 250, .1);padding: 0 2px;border-radius: 2px;height: 21px;line-height: 22px;">torch.fx實現(xiàn)的后量化模型推理速度相比float類型的模型要高3.3倍。并且基于torch.fx實現(xiàn)量化操作相比基于TorchScript IR要簡單很多。

torch.fx還可以做Op融合，F(xiàn)igure7展示了基于torch.fx做了Conv+BN融合后應用在ResNet50上，在n NVIDIA Tesla V100-SXM2 16GB with CUDA version 11.0 和 Intel Xeon Gold 6138 CPU @ 2.00GHz的性能表現(xiàn)，可以看到在GPU上減少了約6%的latency，在CPU上減少了約40%的latency（多線程）和約18%的latency（單線程）。

除此之外torch.fx還可以應用在FLOPs計算，內存帶寬使用分析，工作負載的數(shù)據(jù)值大小估計等，用來分析程序運行時的內存和速度。torch.fx還可以用在形狀推斷，以及模型對應的DAG可視化作圖等等。

最后，torch.fx在runtime階段還支持通過ASIC加速（即將torch.fx中的算子lowering到對應的ASIC上），下面的Figure8展示了基于torch.fx推理ResNet50和LearningToPaint并將算子lowering到TensorRT之后的加速情況：

0x8. 評論

torch.fx對于PyTorch來說確實是一個比較好的工作，因為它消除了一些動態(tài)圖和靜態(tài)圖的Gap。比如在圖改寫方面，torch.fx讓PyTorch想做一些其它靜態(tài)圖框架的算子融合優(yōu)化非常容易。并且torch.fx讓后訓練量化和感知訓練量化以及AMP等的實現(xiàn)難度大大降低，這得益于我們可以直接在Python層操作這個IR，所以我認為這是一個不錯的工作。尤其是對使用PyTorch開發(fā)的算法工程師來說，現(xiàn)在可以基于這個特性大開腦洞了。我之前圍繞FX也做了一個QAT的工作，感興趣可以閱讀：基于OneFlow實現(xiàn)量化感知訓練。

最后總結一下，torch.fx的賣點就是，它使用純Python語言實現(xiàn)了一個可以捕獲PyTorch程序的計算圖并轉化為一個IR的庫，并且非常方便的在這個IR上做Pass，同時提供將變換后的IR Codegen合法的Python代碼功能。我覺得算是達到了在Eager下寫Pass就像做鏈表插入刪除題目一樣順滑。

沐神的論文閱讀方法，感覺確實比較科學，文章末尾再贊一次。