機器之心發布

機器之心編輯部

要想煉丹爽得飛起，就要選擇一個順手的爐子。作為 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 的模型也非常方便，只需數行代碼即可完成：

import torchvision.models as models_torchimport flowvision.models as models_flowresnet101_torch = models_torch.resnet101(pretrained=True)resnet101_flow = models_flow.resnet101()state_dict_torch = resnet101_torch.state_dict()state_dict_numpy = {key: value.detach().cpu().numpy() for key, value in state_dict_torch.items()}resnet101_flow.load_state_dict(state_dict_numpy)

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 函數中構圖。

class ResNet101Graph(oneflow.nn.Graph): def __init__(self, input_shape, input_dtype=oneflow.float32): super().__init__() # 添加 ResNet101 nn.Module self.model = ResNet101Module(input_shape, input_dtype) self.loss_fn = ResNet101_loss_fn # 添加 對應的 Optimizer of_sgd = torch.optim.SGD(self.model.parameters(), lr=1.0, momentum=0.0) self.add_optimizer(of_sgd) # 配置靜態圖的自動優化選項 _config_graph(self) def build(self, input): # 類似 nn.Module 的 forward 方法，這里是構圖，包括了構建后向圖，所以叫 build out = self.model(input) loss = self.loss_fn(out) # build 里面支持構建后向圖 loss.backward() return loss

2）實例化靜態圖：按普通的 Python Class 使用習慣去做初始化就好。

resnet101_graph = ResNet101Graph((args.batch_size, 3, img_shape[1], img_shape[0]))

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

for i in range(m): loss = resnet101_graph(images)

把 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%，一般來說，已經沒有太大優化空間了。但是，OneFlow nn.Graph 下還有一些自動的編譯優化技術可以嘗試。

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

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

def _config_graph(graph): if args.fp16: # 打開 nn.Graph 的自動混合精度執行 graph.config.enable_amp(True) if args.conv_try_run: # 打開 nn.Graph 的卷積的試跑優化 graph.config.enable_cudnn_conv_heuristic_search_algo(False) if args.fuse_add_to_output: # 打開 nn.Graph 的add算子的融合 graph.config.allow_fuse_add_to_output(True) if args.fuse_pad_to_conv: # 打開 nn.Graph 的pad算子的融合 graph.config.allow_fuse_pad_to_conv(True)

對于 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 對不同的參數配置對應的最佳算法進行了預先定義，然后每次搜索時進行匹配得到結果。試運行搜索會傳入實際的張量進行多次試運行，然后返回運行結果。搜索算法返回的結果都是不同算法的元信息及其所需耗時。

啟發式搜索在搜索階段不需額外分配內存，且能更快得到結果；而試運行搜索能得到更為全面和精確的結果，也即通常能更精確地找到最佳算法。啟發式搜索在常見情形下可以得到與試運行搜索一致的結果，但在一些特殊參數配置下無法得到最佳結果。OneFlow 中默認啟動了啟發式搜索，但可通過 graph.config.enable_cudnn_conv_heuristic_search_algo(False) 接口關閉，關閉后使用的就是試運行搜索。

再打開 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 下打開開關就可以，列出來作為參考：

def _config_graph(graph): if args.tensorrt: # 使用 TensorRT 后端執行 graph.config.enable_tensorrt(True)

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.htmlnn.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：

class ResNet101InferenceGraph(oneflow.nn.Graph): def __init__(self): super().__init__() self.model = resnet101_graph.model def build(self, input): return self.model(input) inference_graph = ResNet101InferenceGraph()

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

unused_output = inference_graph(flow.zeros(1, 3, 224, 224))

接下來就可以運行 flow.save 將 inference_graph 的計算圖結構以及權重均保存在 "model" 文件夾下，以供部署使用：

flow.save(inference_graph, "model")

然后只需要運行

docker run --rm --runtime=nvidia --network=host -v$(pwd)/model:/models/resnet101/1 \ oneflowinc/oneflow-serving:nightly

由此可以啟動一個部署著 ResNet101 模型的 Docker 容器。這里的 -v 很重要，它表示將當前目錄下的 model 文件夾映射到容器內的 "/models/resnet101/1" 目錄，其中 /models 是 Triton 讀取模型的默認目錄，Triton 會以該目錄下的一級目錄名（"resnet101"）作為模型名稱，二級目錄名（"1"）作為模型版本。

如果將啟動命令調整為

docker run --rm --runtime=nvidia --network=host -v$(pwd)/model:/models/resnet101/1 \ oneflowinc/oneflow-serving:nightly oneflow-serving --model-store /models --enable-tensorrt resnet101

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

啟動 Docker 容器后，運行下面的命令，就可以查看服務狀態：

curl -v localhost:8000/v2/health/ready

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

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

#/usr/bin/env python3 import numpy as npimport tritonclient.http as httpclientfrom PIL import Image triton_client = httpclient.InferenceServerClient(url='127.0.0.1:8000') image = Image.open("image.jpg")image = image.resize((224, 224))image = np.asarray(image)image = image / 255image = np.expand_dims(image, axis=0)# Transpose NHWC to NCHWimage = np.transpose(image, axes=[0, 3, 1, 2])image = image.astype(np.float32) input = httpclient.InferInput('INPUT_0', image.shape, "FP32")input.set_data_from_numpy(image, binary_data=True)output_placeholder = httpclient.InferRequestedOutput('OUTPUT_0', binary_data=True, class_count=1)output = triton_client.infer("resnet101", inputs=[input], outputs=[output_placeholder]).as_numpy('OUTPUT_0')print(output)

試著運行一下，可以發現它成功的打印出了推理結果：

$ python3 triton_client.py[b'3.630257:499'] # class id 為 499，值為 3.630257

寫在最后

在上述案例中，用戶因時間緊迫沒法做充分調研，抱著試試看的想法選擇了 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