相容PyTorch，25倍性能加速，國産架構OneFlow“超速”了

機器之心釋出

機器之心編輯部

要想煉丹爽得飛起，就要選擇一個順手的爐子。作為 AI 工程師日常必不可缺的「煉丹爐」，「PyTorch 還是 TensorFlow？」已成為知乎、Reddit 等煉丹師出沒之地每年都會讨論的熱門話題。

業界流傳一種說法：PyTorch 适合學術界，TensorFlow 适合工業界。畢竟，PyTorch 是使用者最喜歡的架構，API 非常友好，Eager 模式讓模型搭建和調試過程變得更加容易，不過，它的靜态圖編譯和部署體驗還不令人滿意。TensorFlow 恰恰相反，靜态編譯和部署功能很完備，不過其調試體驗讓人欲哭無淚。

那麼問題來了：魚和熊掌真的不可兼得嗎？未必，來自北京的一流科技團隊推出的開源深度學習架構 OneFlow 已經做到了。

等等，OneFlow 一直主打分布式和高性能，易用性也能和 PyTorch一樣嗎？聽說過 OneFlow 的人一定會發出這樣的疑問。

沒錯，從 2016 年底立項之日起，OneFlow 就是為大規模分布式而生，特色之一就是靜态圖機制，2020 年 7 月在 GitHub 上開源時還不支援動态圖。不過，OneFlow 團隊用一年多時間自研了動态圖引擎， OneFlow 0.7 版本已支援和 PyTorch 一模一樣的 Eager 體驗，也就是說，OneFlow 實作了同時支援動态圖和靜态圖。不僅如此，OneFlow 程式設計 API 完全和 PyTorch 相容，常見深度學習模型隻需修改一行 import oneflow as torch 就可以把 PyTorch 寫的模型在 OneFlow 上跑起來。

不妨先到 OneFlow 視覺模型庫 flowvision 看一看：https://github.com/Oneflow-Inc/vision ，這個模型庫已經支援計算機視覺領域圖像分類、分割和檢測等方向的經典 SOTA 模型 （見下表），這些模型都可以通過 import torch as flow 或 import oneflow as torch 實作自由切換。

OneFlow 和 PyTorch 相容之後，使用者可以像使用 PyTorch 一樣來使用 OneFlow ，對模型效果比較滿意之後，可以繼續使用 OneFlow 擴充到大規模分布式或使用靜态圖部署模型。聽上去是不是 too good to be true？

在下面的案例中，一家頭部通信公司基于 PyTorch 的業務模型快速友善地遷移成 OneFlow 的模型，并進行大幅度的訓練/推理性能優化、部署上線，短短幾天時間就讓業務得以按時上線部署，且各項性能名額均大幅超出預期！

他們究竟是是如何做到的？先從項目背景說起。

為什麼選擇 OneFlow？

因業務發展需求，這家通信公司近期将上線一款基于深度學習的圖像識别應用，該項目的業務需求有如下五個特點：

資料量大：資料庫中有過億級别的圖檔

模型簡單：比較正常的分類模型

400 多張顯示卡，短期内無法擴容

對于訓練/推理的吞吐有硬性名額

上線時間緊迫

使用者基于市面上最流行的深度學習架構 PyTorch 搭建了業務模型，且跑通了正常訓練流程，但是訓練/推理都很慢，遠遠達不到目标（離上線 QPS 有 20 倍的差距），随着傳遞日期臨近，整個團隊深陷焦慮。

使用者嘗試了各種方案（基于已有實作進行優化）都無濟于事，于是調研了其他深度學習架構，如 TensorFlow、OneFlow 等，發現 OneFlow (https://github.com/OneFlow-Inc/oneflow) 是加速 PyTorch 風格代碼的最平滑架構。

具體而言，使用者選擇試用 OneFlow的理由主要有三點：

1、OneFlow 是衆多深度學習架構中，API 與 PyTorch 相容性最高的，這樣友善工程師用最少的時間/人力成本，對已有項目代碼進行遷移，減少學習成本。

2、OneFlow 動靜轉換十分友善，動态圖（Eager）模式的代碼簡單改動幾行就能轉換為靜态圖（nn.Graph）模式。

3、OneFlow 在架構層面做了大量優化，nn.Graph 提供了簡潔、豐富的性能優化選項，如算子融合（Kernel Fusion）、自動混合精度訓練 (Auto Mixed Precision Training) 等。

于是，使用者就開始嘗試将已有代碼遷移至 OneFlow，沒想到，不到半天就搞定并跑起來了，遷移過程非常絲滑。

在 OneFlow 官方文檔(https://docs.oneflow.org/master/index.html) 以及 OneFlow 研發團隊的大力支援下，使用者開展了以下工作：

将已有 PyTorch 的項目代碼完全遷移到 OneFlow

将項目代碼由動态圖模式（Eager Mode）改造為靜态圖模式（Graph Mode）

開啟 OneFlow Graph 模式下的各種優化選項并訓練模型

用 Serving 子產品部署模型上線

遷移調優過程

1. 一鍵遷移 PyTorch 模型轉OneFlow模型：隻需import oneflow as torch就夠了

OneFlow 最新釋出的 0.7.0 版本對 PyTorch 接口的相容性有了進一步的完善。OneFlow 對已經支援的算子都能保證和 PyTorch 的接口在語義和結果上一緻。于是使用者就嘗試了一下遷移模型腳本到 OneFlow。由于業務模型的主幹網絡是 resnet101，在遷移過程中，使用者參考了官方文檔(https://docs.oneflow.org/master/cookies/torch2flow.html)來遷移 ，發現隻需要模型檔案中與 torch 相關的 import 修改為import oneflow as torch，就完成了模型代碼的遷移工作。

在模型腳本遷移完畢之後，還需要驗證模型遷移的正确性，看看精度是不是對齊了。

1）使用者首先做了推理精度的驗證，就是直接加載 PyTorch 訓練好的模型然後驗證推理精度，由于 OneFlow 對齊了 PyTorch 的接口，是以加載 PyTorch 的模型也非常友善，隻需數行代碼即可完成：

2）在驗證完推理精度後接着就是驗證訓練流程，在對齊訓練超參數之後，使用 OneFlow 訓練模型的 loss 曲線和 PyTorch 的收斂曲線也一緻，在小資料集上的精度完全一緻。

2. 使用 OneFlow 的 nn.Graph 加速模型訓練與推理性能

在驗證完算法正确性後，就需要考慮如何加速執行了。如果使用現有的動态圖模型直接部署，在現有的機器資源和時間限制内，使用最原始的代碼實作還差約 20 倍的性能，短期内是一個不可能完成的任務。

使用者決定雙管齊下，在基于 PyTorch 做加速優化時，并行地使用 OneFlow 進行加速。最終結合「動态轉靜态、算法邏輯約減、提高并行度、靜态編譯優化」這四類技巧，最終單機執行達到了 25 倍以上的加速效果。

2.1 動态轉靜态

動态圖轉靜态圖執行後，得到了約 25% 的性能加速。

OneFlow 有個 ResNet50 的開源項目（ https://github.com/Oneflow-Inc/models/tree/main/Vision/classification/image/resnet50 ），了解到單卡的執行效率已經做得很高，照貓畫虎，這些優化技巧都可以用在 ResNet101 上。

OneFlow ResNet50 下做模型加速使用的是靜态圖nn.Graph，類似 PyTorch 的 TorchScript。但OneFlow的優化功能做的更全面一些，運作時也是一個特有的服務于加速的 Actor Runtime。

nn.Graph 是一個面向對象風格的靜态圖類，它代表一個完整的靜态計算圖。對于預測任務，nn.Graph 可以隻包括前向計算；對于訓練任務，還可以包括後向計算和模型更新。

nn.Graph 的基礎接口和nn.Module 的行為比較類似，比如添加子 Module，自定義算法執行邏輯，調用以執行一次計算，儲存模型等。被添加進入nn.Graph 的nn.Module對象，在nn.Graph裡執行時，就會采用靜态圖模式執行，如此動态圖下的計算邏輯就可以被靜态圖直接複用，這樣就實作了動靜執行的切換。特殊一點的是，Optimizer 也可以添加進入靜态圖，這樣前向、後向、模型更新可以被加入一個完整的靜态圖做聯合優化。

下面的步驟把動态執行的 ResNet101Module 變成靜态執行，使用方式和nn.Module 類似，隻需要聲明、執行個體化、調用三個基本步驟。

1）聲明一個靜态圖：主要包括兩部分，先在初始化函數中添加要靜态化的nn.Module 和Optimizer；然後在build 函數中構圖。

2）執行個體化靜态圖：按普通的 Python Class 使用習慣去做初始化就好。

3）調用靜态圖：類似nn.Module 的調用方式，注意第一次調用會觸發編譯，是以第一次調用比後面的時間要長。

把 ResNet101 的nn.Module 的執行個體加入nn.Graph 執行後，對比得到約 25% 的加速。

2.2 算法層次的優化

使用者在把動态圖代碼遷移到靜态圖代碼的過程中，因為需要考慮哪些部分要做靜态化，是以對模型做了子產品化的重構，但發現本任務中有些計算是做實驗時遺留的，在部署時并不必要，順便做了算法邏輯的約減：

一般推理時隻需要前向計算，後向計算是不需要的，但在使用者這個特殊的模型裡，部署和推理也是需要後向計算，隻是不需要模型更新，這就導緻使用者寫代碼時為了保留後向計算也誤把參數更新的邏輯保留下來了。據此可以省略參數的梯度計算，這裡大概帶來了 75% 的加速；

進而發現原任務(前向、後向、前向）中的第二次前向在部署時是多餘的，可以裁剪掉，這裡大概帶來了大約 33% 的加速。

總體而言，算法層次方面累積加速了 2.33 倍，事實證明，算法邏輯本身具有很大的優化空間，代碼做好子產品化，可以比較容易找到算法邏輯上的優化點。當然，這部分改善也适用于PyTorch。

2.3 提高并行度

這個思路也比較直接，在做完優化的基礎上，使用者觀察到 GPU 的使用率隻有 30%。此時 batch_size 為 1（ BN 的某些參數和 batch 大小有關，原先使用者擔心擴大 batch_size 可能影響計算結果，事後證明這個擔心是多餘的，從理論推導和實驗結果都證明，擴大 batch_size 并不影響計算結果），單程序，提高資料并行度是很值得嘗試的方案。是以，使用者嘗試了提高 batch_size 和 多程序方案：

增大 batch_size，預設 batch_size 為 1，此時 GPU 使用率為 30%，當增大到 16 時，最高可以達到 90%，這裡大約得到了 155% 的加速；

由于資料預處理在 CPU，網絡計算在 GPU，兩種裝置接力執行，這時使用 2 程序進行，給資料加載部分加一個互斥鎖，可以比較簡易的實作 CPU 和 GPU 兩級流水線，這裡帶來了 80% 的加速。

提高并行度的累積加速是 4.6 倍。增加并行度以充分利用多核、多裝置，帶來了最明顯的加速效果。當然，這裡的優化效果是使用者遷移到 OneFlow 後實作的，在 PyTorch 上也可以做到。

2.4 靜态編譯優化

做到以上優化後，GPU 使用率已經能比較穩定的保持在 90%，一般來說，已經沒有太大優化空間了。但是，OneFlownn.Graph 下還有一些自動的編譯優化技術可以嘗試。

比如利用自動混合精度做低精度計算、利用算子融合來減少訪存開銷等，這裡最終帶來了 64% 的加速，速度到了原來最好性能的 1.56 倍。

此前示例中提到的_config_graph 函數就是在配置這些優化選項，具體如下：

對于 ResNet101，batch_size 設定為 16，在nn.Graph 無優化選項打開的基礎上：

打開混合精度，測試得到了 36% 的加速

自動混合精度訓練，自動将網絡中的合适的算子由 FP32 單精度計算轉換成 FP16 半精度浮點進行計算，不僅可以減少 GPU 顯存占用，而且可以提升整體性能，在支援 Tensor Core 的 GPU 裝置上還會使用 Tensor Core 進一步加速訓練。

再打開卷積試跑優化，測試得到了 7% 的加速，總加速為 43%

cudnn 的 convolution 算子包含多種算法，例如前向的算法(https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnConvolutionFwdAlgo_t)。不同的 input 和 filter 大小在不同的算法下有不同的性能表現，為了選擇最佳算法，在調用 cudnn convolution 算子接口前，需要先調用 cudnn convolution searching algorithm 的接口。cudnn 提供了2種搜尋模式：啟發式搜尋(https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnGetConvolutionForwardAlgorithm_v7)和試運作搜尋(cudnnFindConvolutionForwardAlgorithm)(https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnFindConvolutionForwardAlgorithm)。

啟發式搜尋是通過一種「查表」的方式來搜尋最佳算法，cudnn 對不同的參數配置對應的最佳算法進行了預先定義，然後每次搜尋時進行比對得到結果。試運作搜尋會傳入實際的張量進行多次試運作，然後傳回運作結果。搜尋算法傳回的結果都是不同算法的元資訊及其所需耗時。

再打開 pad 和 conv 算子融合，測試得到了 19% 的加速，總加速為 62%

在 CNN 網絡 Backbone 中有很多 convolution + pad 的組合，convolution 算子自身支援 pad 操作，自動将 pad 算子 fuse 到 convolution 算子上，可以省掉 pad 算子的開銷，提升網絡整體性能。

再打開 add 的算子的融合，測試得到了 2% 的加速，總加速為 64%

自動将網絡中常見的訪存密集型算子 Elementwise add 算子和上遊的算子 fuse 起來，可以減少帶寬使用，進而提升性能。對于 Elementwise add 算子來說，将其 fuse 到上一個算子，可以減少一次資料讀寫，有約 2/3 的性能提升。

另外nn.Graph 可以很友善地支援使用 TensorRT 。本優化對象沒有更新模型的需求，是以也适合使用 TensorRT 做加速。在nn.Graph 無優化選項基礎上， batch_size 設定為 16，新增自動混合精度、NHWC、使用 TensorRT 後端，可以提速 48%。

在這個模型裡，隻使用 TensorRT 後端比隻使用 OneFlow 的靜态圖優化還差一點，可能的原因是， TensorRT 下的一些優化在nn.Graph裡已經做了，是以沒有帶來額外收益。不過其實驗起來還比較友善，編譯一下帶 TensorRT 的 OneFlow，再在nn.Graph 下打開開關就可以，列出來作為參考：

2.5 加速優化總結

以上記錄了加速的主要過程，動态轉靜态加速約 1.25 倍、算法邏輯約減加速約 2.33 倍、提高并行度加速約 4.6 倍、靜态編譯優化加速約 1.6 倍，累積加速約 21 倍。中間有些小的優化點沒有完全記錄，實際累積的加速效果達到了 25 倍以上，超過了項目部署的 20 倍加速需求。

nn.Graph 的進一步的使用可以參考：

nn.Graph 的使用教程，https://docs.oneflow.org/en/master/basics/08_nn_graph.html

nn.Graph 的 API 文檔，https://oneflow.readthedocs.io/en/master/graph.html

3. 使用 OneFlow-Serving，輕松将訓練好的模型部署上線

當使用者完成訓練，得到最終的模型之後，接下來的一步就是模型部署。不同于模型訓練時需要進行權重更新，部署時的權重固定不變，是以可以進行更激進的速度優化，例如 int8 量化、更廣泛的 kernel fusion、constant folding 等等。

使用者參考 OneFlow v0.7.0 提供了官方的 Serving 子產品(https://github.com/Oneflow-Inc/serving)，它是一個 NVIDIA Triton 的後端，內建了 OneFlow 内置的 XRT 子產品，并提供了開箱即用的使用者接口。隻需使用下述方法就将訓練好的 OneFlow 模型快速高效的部署起來：

為了将模型用于推理，在使用 nn.Graph 訓練完成之後，需要構造一個隻包含前向的ResNet101InferenceGraph：

并以一個樣例輸入運作 inference_graph，觸發 inference_graph 的計算圖建構：

接下來就可以運作flow.save 将 inference_graph 的計算圖結構以及權重均儲存在 "model" 檔案夾下，以供部署使用：

然後隻需要運作

由此可以啟動一個部署着 ResNet101 模型的 Docker 容器。這裡的 -v 很重要，它表示将目前目錄下的 model 檔案夾映射到容器内的 "/models/resnet101/1" 目錄，其中 /models 是 Triton 讀取模型的預設目錄，Triton 會以該目錄下的一級目錄名（"resnet101"）作為模型名稱，二級目錄名（"1"）作為模型版本。

如果将啟動指令調整為

模型就會通過 OneFlow 的 XRT 子產品自動使用 TensorRT 進行推理，此外 OneFlow Serving 還支援類似的 “--enable-openvino”。

啟動 Docker 容器後，運作下面的指令，就可以檢視服務狀态：

傳回值為HTTP/1.1 200 OK，表示服務正在正常工作。

接下來就可以使用 Triton 的 C++ 或 Python SDK 實作向服務端發送請求并擷取結果的邏輯了，例如一個最簡單的用戶端：

試着運作一下，可以發現它成功的列印出了推理結果：

寫在最後

在上述案例中，使用者因時間緊迫沒法做充分調研，抱着試試看的想法選擇了 OneFlow，幸運的是，終于在極限壓縮的項目周期裡順利完成了任務。

基于 OneFlow v0.7.0 ，使用者輕松地将之前開發的 PyTorch 的業務模型代碼一鍵遷移成 OneFlow 的模型代碼，再經過簡單加工就轉成 OneFlow 的靜态圖 nn.Graph 模式，并利用 nn.Graph 豐富、高效、簡潔的優化開關來快速大幅提升模型的訓練速度，利用完善的周邊工具鍊如 OneFlow-Serving 友善的進行線上部署。值得一提的是，使用者還可以使用 OneFlow-ONNX 工具将 OneFlow 高效訓練好的模型轉成 ONNX 格式導入到其他架構中使用。

本文隻介紹了借助和 PyTorch 的相容性 OneFlow 幫助使用者實作模型加速和部署的例子。OneFlow 原來的殺手锏功能“大規模分布式”還沒有展現出來，未來，我們将進一步介紹 OneFlow 如何幫助習慣 PyTorch 的使用者便捷地實作大規模預訓練 Transformer 模型和搜尋推薦廣告領域需要的大規模 embedding 模型。

OneFlow項目位址：https://github.com/Oneflow-Inc/oneflow/

OneFlow使用者文檔：https://docs.oneflow.org/master/index.html