機(jī)器之心發(fā)布

機(jī)器之心編輯部

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

業(yè)界流傳一種說法：PyTorch 適合學(xué)術(shù)界，TensorFlow 適合工業(yè)界。畢竟，PyTorch 是用戶最喜歡的框架，API 非常友好，Eager 模式讓模型搭建和調(diào)試過程變得更加容易，不過，它的靜態(tài)圖編譯和部署體驗(yàn)還不令人滿意。TensorFlow 恰恰相反，靜態(tài)編譯和部署功能很完備，不過其調(diào)試體驗(yàn)讓人欲哭無淚。

那么問題來了：魚和熊掌真的不可兼得嗎？未必，來自北京的一流科技團(tuán)隊(duì)推出的開源深度學(xué)習(xí)框架 OneFlow 已經(jīng)做到了。

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

沒錯(cuò)，從 2016 年底立項(xiàng)之日起，OneFlow 就是為大規(guī)模分布式而生，特色之一就是靜態(tài)圖機(jī)制，2020 年 7 月在 GitHub 上開源時(shí)還不支持動(dòng)態(tài)圖。不過，OneFlow 團(tuán)隊(duì)用一年多時(shí)間自研了動(dòng)態(tài)圖引擎， OneFlow 0.7 版本已支持和 PyTorch 一模一樣的 Eager 體驗(yàn)，也就是說，OneFlow 實(shí)現(xiàn)了同時(shí)支持動(dòng)態(tài)圖和靜態(tài)圖。不僅如此，OneFlow 編程 API 完全和 PyTorch 兼容，常見深度學(xué)習(xí)模型只需修改一行 import oneflow as torch 就可以把 PyTorch 寫的模型在 OneFlow 上跑起來。

不妨先到 OneFlow 視覺模型庫 flowvision 看一看：https://github.com/Oneflow-Inc/vision ，這個(gè)模型庫已經(jīng)支持計(jì)算機(jī)視覺領(lǐng)域圖像分類、分割和檢測(cè)等方向的經(jīng)典 SOTA 模型 （見下表），這些模型都可以通過 import torch as flow 或 import oneflow as torch 實(shí)現(xiàn)自由切換。

OneFlow 和 PyTorch 兼容之后，用戶可以像使用 PyTorch 一樣來使用 OneFlow ，對(duì)模型效果比較滿意之后，可以繼續(xù)使用 OneFlow 擴(kuò)展到大規(guī)模分布式或使用靜態(tài)圖部署模型。聽上去是不是 too good to be true？

在下面的案例中，一家頭部通信公司基于 PyTorch 的業(yè)務(wù)模型快速方便地遷移成 OneFlow 的模型，并進(jìn)行大幅度的訓(xùn)練/推理性能優(yōu)化、部署上線，短短幾天時(shí)間就讓業(yè)務(wù)得以按時(shí)上線部署，且各項(xiàng)性能指標(biāo)均大幅超出預(yù)期！

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

為什么選擇 OneFlow？

因業(yè)務(wù)發(fā)展需求，這家通信公司近期將上線一款基于深度學(xué)習(xí)的圖像識(shí)別應(yīng)用，該項(xiàng)目的業(yè)務(wù)需求有如下五個(gè)特點(diǎn)：

數(shù)據(jù)量大：數(shù)據(jù)庫中有過億級(jí)別的圖片模型簡(jiǎn)單：比較常規(guī)的分類模型400 多張顯卡，短期內(nèi)無法擴(kuò)容對(duì)于訓(xùn)練/推理的吞吐有硬性指標(biāo)上線時(shí)間緊迫

用戶基于市面上最流行的深度學(xué)習(xí)框架 PyTorch 搭建了業(yè)務(wù)模型，且跑通了正常訓(xùn)練流程，但是訓(xùn)練/推理都很慢，遠(yuǎn)遠(yuǎn)達(dá)不到目標(biāo)（離上線 QPS 有 20 倍的差距），隨著交付日期臨近，整個(gè)團(tuán)隊(duì)深陷焦慮。

用戶嘗試了各種方案（基于已有實(shí)現(xiàn)進(jìn)行優(yōu)化）都無濟(jì)于事，于是調(diào)研了其他深度學(xué)習(xí)框架，如 TensorFlow、OneFlow 等，發(fā)現(xiàn) OneFlow (https://github.com/OneFlow-Inc/oneflow) 是加速 PyTorch 風(fēng)格代碼的最平滑框架。

具體而言，用戶選擇試用 OneFlow的理由主要有三點(diǎn)：

1、OneFlow 是眾多深度學(xué)習(xí)框架中，API 與 PyTorch 兼容性最高的，這樣方便工程師用最少的時(shí)間/人力成本，對(duì)已有項(xiàng)目代碼進(jìn)行遷移，減少學(xué)習(xí)成本。

2、OneFlow 動(dòng)靜轉(zhuǎn)換十分方便，動(dòng)態(tài)圖（Eager）模式的代碼簡(jiǎn)單改動(dòng)幾行就能轉(zhuǎn)換為靜態(tài)圖（nn.Graph）模式。

3、OneFlow 在框架層面做了大量?jī)?yōu)化，nn.Graph 提供了簡(jiǎn)潔、豐富的性能優(yōu)化選項(xiàng)，如算子融合（Kernel Fusion）、自動(dòng)混合精度訓(xùn)練 (Auto Mixed Precision Training) 等。

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

在 OneFlow 官方文檔(https://docs.oneflow.org/master/index.html) 以及 OneFlow 研發(fā)團(tuán)隊(duì)的大力支持下，用戶開展了以下工作：

將已有 PyTorch 的項(xiàng)目代碼完全遷移到 OneFlow將項(xiàng)目代碼由動(dòng)態(tài)圖模式（Eager Mode）改造為靜態(tài)圖模式（Graph Mode）開啟 OneFlow Graph 模式下的各種優(yōu)化選項(xiàng)并訓(xùn)練模型用 Serving 模塊部署模型上線

遷移調(diào)優(yōu)過程

1. 一鍵遷移 PyTorch 模型轉(zhuǎn) OneFlow 模型：只需 import oneflow as torch 就夠了

OneFlow 最新發(fā)布的 0.7.0 版本對(duì) PyTorch 接口的兼容性有了進(jìn)一步的完善。OneFlow 對(duì)已經(jīng)支持的算子都能保證和 PyTorch 的接口在語義和結(jié)果上一致。于是用戶就嘗試了一下遷移模型腳本到 OneFlow。由于業(yè)務(wù)模型的主干網(wǎng)絡(luò)是 resnet101，在遷移過程中，用戶參考了官方文檔(https://docs.oneflow.org/master/cookies/torch2flow.html)來遷移 ，發(fā)現(xiàn)只需要模型文件中與 torch 相關(guān)的 import 修改為 import oneflow as torch，就完成了模型代碼的遷移工作。

在模型腳本遷移完畢之后，還需要驗(yàn)證模型遷移的正確性，看看精度是不是對(duì)齊了。

1）用戶首先做了推理精度的驗(yàn)證，就是直接加載 PyTorch 訓(xùn)練好的模型然后驗(yàn)證推理精度，由于 OneFlow 對(duì)齊了 PyTorch 的接口，所以加載 PyTorch 的模型也非常方便，只需數(shù)行代碼即可完成：

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）在驗(yàn)證完推理精度后接著就是驗(yàn)證訓(xùn)練流程，在對(duì)齊訓(xùn)練超參數(shù)之后，使用 OneFlow 訓(xùn)練模型的 loss 曲線和 PyTorch 的收斂曲線也一致，在小數(shù)據(jù)集上的精度完全一致。

2. 使用 OneFlow 的 nn.Graph 加速模型訓(xùn)練與推理性能

在驗(yàn)證完算法正確性后，就需要考慮如何加速執(zhí)行了。如果使用現(xiàn)有的動(dòng)態(tài)圖模型直接部署，在現(xiàn)有的機(jī)器資源和時(shí)間限制內(nèi)，使用最原始的代碼實(shí)現(xiàn)還差約 20 倍的性能，短期內(nèi)是一個(gè)不可能完成的任務(wù)。

用戶決定雙管齊下，在基于 PyTorch 做加速優(yōu)化時(shí)，并行地使用 OneFlow 進(jìn)行加速。最終結(jié)合「動(dòng)態(tài)轉(zhuǎn)靜態(tài)、算法邏輯約減、提高并行度、靜態(tài)編譯優(yōu)化」這四類技巧，最終單機(jī)執(zhí)行達(dá)到了 25 倍以上的加速效果。

2.1 動(dòng)態(tài)轉(zhuǎn)靜態(tài)

動(dòng)態(tài)圖轉(zhuǎn)靜態(tài)圖執(zhí)行后，得到了約 25% 的性能加速。

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

OneFlow ResNet50 下做模型加速使用的是靜態(tài)圖 nn.Graph，類似 PyTorch 的 TorchScript。但OneFlow的優(yōu)化功能做的更全面一些，運(yùn)行時(shí)也是一個(gè)特有的服務(wù)于加速的 Actor Runtime。

nn.Graph 是一個(gè)面向?qū)ο箫L(fēng)格的靜態(tài)圖類，它代表一個(gè)完整的靜態(tài)計(jì)算圖。對(duì)于預(yù)測(cè)任務(wù)，nn.Graph 可以只包括前向計(jì)算；對(duì)于訓(xùn)練任務(wù)，還可以包括后向計(jì)算和模型更新。

nn.Graph 的基礎(chǔ)接口和 nn.Module 的行為比較類似，比如添加子 Module，自定義算法執(zhí)行邏輯，調(diào)用以執(zhí)行一次計(jì)算，保存模型等。被添加進(jìn)入 nn.Graph 的 nn.Module 對(duì)象，在 nn.Graph 里執(zhí)行時(shí)，就會(huì)采用靜態(tài)圖模式執(zhí)行，如此動(dòng)態(tài)圖下的計(jì)算邏輯就可以被靜態(tài)圖直接復(fù)用，這樣就實(shí)現(xiàn)了動(dòng)靜執(zhí)行的切換。特殊一點(diǎn)的是，Optimizer 也可以添加進(jìn)入靜態(tài)圖，這樣前向、后向、模型更新可以被加入一個(gè)完整的靜態(tài)圖做聯(lián)合優(yōu)化。

下面的步驟把動(dòng)態(tài)執(zhí)行的 ResNet101Module 變成靜態(tài)執(zhí)行，使用方式和 nn.Module 類似，只需要聲明、實(shí)例化、調(diào)用三個(gè)基本步驟。

1）聲明一個(gè)靜態(tài)圖：主要包括兩部分，先在初始化函數(shù)中添加要靜態(tài)化的 nn.Module 和 Optimizer；然后在 build 函數(shù)中構(gòu)圖。

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 # 添加 對(duì)應(yīng)的 Optimizer of_sgd = torch.optim.SGD(self.model.parameters(), lr=1.0, momentum=0.0) self.add_optimizer(of_sgd) # 配置靜態(tài)圖的自動(dòng)優(yōu)化選項(xiàng) _config_graph(self) def build(self, input): # 類似 nn.Module 的 forward 方法，這里是構(gòu)圖，包括了構(gòu)建后向圖，所以叫 build out = self.model(input) loss = self.loss_fn(out) # build 里面支持構(gòu)建后向圖 loss.backward() return loss

2）實(shí)例化靜態(tài)圖：按普通的 Python Class 使用習(xí)慣去做初始化就好。

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

3）調(diào)用靜態(tài)圖：類似 nn.Module 的調(diào)用方式，注意第一次調(diào)用會(huì)觸發(fā)編譯，所以第一次調(diào)用比后面的時(shí)間要長(zhǎng)。

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

把 ResNet101 的 nn.Module 的實(shí)例加入 nn.Graph 執(zhí)行后，對(duì)比得到約 25% 的加速。

2.2 算法層次的優(yōu)化

用戶在把動(dòng)態(tài)圖代碼遷移到靜態(tài)圖代碼的過程中，因?yàn)樾枰紤]哪些部分要做靜態(tài)化，所以對(duì)模型做了模塊化的重構(gòu)，但發(fā)現(xiàn)本任務(wù)中有些計(jì)算是做實(shí)驗(yàn)時(shí)遺留的，在部署時(shí)并不必要，順便做了算法邏輯的約減：

一般推理時(shí)只需要前向計(jì)算，后向計(jì)算是不需要的，但在用戶這個(gè)特殊的模型里，部署和推理也是需要后向計(jì)算，只是不需要模型更新，這就導(dǎo)致用戶寫代碼時(shí)為了保留后向計(jì)算也誤把參數(shù)更新的邏輯保留下來了。據(jù)此可以省略參數(shù)的梯度計(jì)算，這里大概帶來了 75% 的加速；進(jìn)而發(fā)現(xiàn)原任務(wù)(前向、后向、前向）中的第二次前向在部署時(shí)是多余的，可以裁剪掉，這里大概帶來了大約 33% 的加速。

總體而言，算法層次方面累積加速了 2.33 倍，事實(shí)證明，算法邏輯本身具有很大的優(yōu)化空間，代碼做好模塊化，可以比較容易找到算法邏輯上的優(yōu)化點(diǎn)。當(dāng)然，這部分改善也適用于PyTorch。

2.3 提高并行度

這個(gè)思路也比較直接，在做完優(yōu)化的基礎(chǔ)上，用戶觀察到 GPU 的利用率只有 30%。此時(shí) batch_size 為 1（ BN 的某些參數(shù)和 batch 大小有關(guān)，原先用戶擔(dān)心擴(kuò)大 batch_size 可能影響計(jì)算結(jié)果，事后證明這個(gè)擔(dān)心是多余的，從理論推導(dǎo)和實(shí)驗(yàn)結(jié)果都證實(shí)，擴(kuò)大 batch_size 并不影響計(jì)算結(jié)果），單進(jìn)程，提高數(shù)據(jù)并行度是很值得嘗試的方案。因此，用戶嘗試了提高 batch_size 和 多進(jìn)程方案：

增大 batch_size，默認(rèn) batch_size 為 1，此時(shí) GPU 利用率為 30%，當(dāng)增大到 16 時(shí)，最高可以達(dá)到 90%，這里大約得到了 155% 的加速；由于數(shù)據(jù)預(yù)處理在 CPU，網(wǎng)絡(luò)計(jì)算在 GPU，兩種設(shè)備接力執(zhí)行，這時(shí)使用 2 進(jìn)程進(jìn)行，給數(shù)據(jù)加載部分加一個(gè)互斥鎖，可以比較簡(jiǎn)易的實(shí)現(xiàn) CPU 和 GPU 兩級(jí)流水線，這里帶來了 80% 的加速。

提高并行度的累積加速是 4.6 倍。增加并行度以充分利用多核、多設(shè)備，帶來了最明顯的加速效果。當(dāng)然，這里的優(yōu)化效果是用戶遷移到 OneFlow 后實(shí)現(xiàn)的，在 PyTorch 上也可以做到。

2.4 靜態(tài)編譯優(yōu)化

做到以上優(yōu)化后，GPU 利用率已經(jīng)能比較穩(wěn)定的保持在 90%，一般來說，已經(jīng)沒有太大優(yōu)化空間了。但是，OneFlow nn.Graph 下還有一些自動(dòng)的編譯優(yōu)化技術(shù)可以嘗試。

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

此前示例中提到的 _config_graph 函數(shù)就是在配置這些優(yōu)化選項(xiàng)，具體如下：

def _config_graph(graph): if args.fp16: # 打開 nn.Graph 的自動(dòng)混合精度執(zhí)行 graph.config.enable_amp(True) if args.conv_try_run: # 打開 nn.Graph 的卷積的試跑優(yōu)化 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)

對(duì)于 ResNet101，batch_size 設(shè)置為 16，在 nn.Graph 無優(yōu)化選項(xiàng)打開的基礎(chǔ)上：

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

自動(dòng)混合精度訓(xùn)練，自動(dòng)將網(wǎng)絡(luò)中的合適的算子由 FP32 單精度計(jì)算轉(zhuǎn)換成 FP16 半精度浮點(diǎn)進(jìn)行計(jì)算，不僅可以減少 GPU 顯存占用，而且可以提升整體性能，在支持 Tensor Core 的 GPU 設(shè)備上還會(huì)使用 Tensor Core 進(jìn)一步加速訓(xùn)練。

再打開卷積試跑優(yōu)化，測(cè)試得到了 7% 的加速，總加速為 43%

cudnn 的 convolution 算子包含多種算法，例如前向的算法(https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnConvolutionFwdAlgo_t)。不同的 input 和 filter 大小在不同的算法下有不同的性能表現(xiàn)，為了選擇最佳算法，在調(diào)用 cudnn convolution 算子接口前，需要先調(diào)用 cudnn convolution searching algorithm 的接口。cudnn 提供了2種搜索模式：?jiǎn)l(fā)式搜索(https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnGetConvolutionForwardAlgorithm_v7)和試運(yùn)行搜索(cudnnFindConvolutionForwardAlgorithm)(https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnFindConvolutionForwardAlgorithm)。

啟發(fā)式搜索是通過一種「查表」的方式來搜尋最佳算法，cudnn 對(duì)不同的參數(shù)配置對(duì)應(yīng)的最佳算法進(jìn)行了預(yù)先定義，然后每次搜索時(shí)進(jìn)行匹配得到結(jié)果。試運(yùn)行搜索會(huì)傳入實(shí)際的張量進(jìn)行多次試運(yùn)行，然后返回運(yùn)行結(jié)果。搜索算法返回的結(jié)果都是不同算法的元信息及其所需耗時(shí)。

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

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

在 CNN 網(wǎng)絡(luò) Backbone 中有很多 convolution + pad 的組合，convolution 算子自身支持 pad 操作，自動(dòng)將 pad 算子 fuse 到 convolution 算子上，可以省掉 pad 算子的開銷，提升網(wǎng)絡(luò)整體性能。

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

自動(dòng)將網(wǎng)絡(luò)中常見的訪存密集型算子 Elementwise add 算子和上游的算子 fuse 起來，可以減少帶寬使用，從而提升性能。對(duì)于 Elementwise add 算子來說，將其 fuse 到上一個(gè)算子，可以減少一次數(shù)據(jù)讀寫，有約 2/3 的性能提升。

另外 nn.Graph 可以很方便地支持使用 TensorRT 。本優(yōu)化對(duì)象沒有更新模型的需求，所以也適合使用 TensorRT 做加速。在 nn.Graph 無優(yōu)化選項(xiàng)基礎(chǔ)上， batch_size 設(shè)置為 16，新增自動(dòng)混合精度、NHWC、使用 TensorRT 后端，可以提速 48%。

在這個(gè)模型里，只使用 TensorRT 后端比只使用 OneFlow 的靜態(tài)圖優(yōu)化還差一點(diǎn)，可能的原因是， TensorRT 下的一些優(yōu)化在 nn.Graph 里已經(jīng)做了，所以沒有帶來額外收益。不過其實(shí)驗(yàn)起來還比較方便，編譯一下帶 TensorRT 的 OneFlow，再在 nn.Graph 下打開開關(guān)就可以，列出來作為參考：

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

2.5 加速優(yōu)化總結(jié)

以上記錄了加速的主要過程，動(dòng)態(tài)轉(zhuǎn)靜態(tài)加速約 1.25 倍、算法邏輯約減加速約 2.33 倍、提高并行度加速約 4.6 倍、靜態(tài)編譯優(yōu)化加速約 1.6 倍，累積加速約 21 倍。中間有些小的優(yōu)化點(diǎn)沒有完全記錄，實(shí)際累積的加速效果達(dá)到了 25 倍以上，超過了項(xiàng)目部署的 20 倍加速需求。

nn.Graph 的進(jìn)一步的使用可以參考：

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，輕松將訓(xùn)練好的模型部署上線

當(dāng)用戶完成訓(xùn)練，得到最終的模型之后，接下來的一步就是模型部署。不同于模型訓(xùn)練時(shí)需要進(jìn)行權(quán)重更新，部署時(shí)的權(quán)重固定不變，所以可以進(jìn)行更激進(jìn)的速度優(yōu)化，例如 int8 量化、更廣泛的 kernel fusion、constant folding 等等。

用戶參考 OneFlow v0.7.0 提供了官方的 Serving 模塊(https://github.com/Oneflow-Inc/serving)，它是一個(gè) NVIDIA Triton 的后端，集成了 OneFlow 內(nèi)置的 XRT 模塊，并提供了開箱即用的用戶接口。只需使用下述方法就將訓(xùn)練好的 OneFlow 模型快速高效的部署起來：

為了將模型用于推理，在使用 nn.Graph 訓(xùn)練完成之后，需要構(gòu)造一個(gè)只包含前向的 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()

并以一個(gè)樣例輸入運(yùn)行 inference_graph，觸發(fā) inference_graph 的計(jì)算圖構(gòu)建：

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

接下來就可以運(yùn)行 flow.save 將 inference_graph 的計(jì)算圖結(jié)構(gòu)以及權(quán)重均保存在 "model" 文件夾下，以供部署使用：

flow.save(inference_graph, "model")

然后只需要運(yùn)行

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

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

如果將啟動(dòng)命令調(diào)整為

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

模型就會(huì)通過 OneFlow 的 XRT 模塊自動(dòng)使用 TensorRT 進(jìn)行推理，此外 OneFlow Serving 還支持類似的 “--enable-openvino”。

啟動(dòng) Docker 容器后，運(yùn)行下面的命令，就可以查看服務(wù)狀態(tài)：

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

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

接下來就可以使用 Triton 的 C++ 或 Python SDK 實(shí)現(xiàn)向服務(wù)端發(fā)送請(qǐng)求并獲取結(jié)果的邏輯了，例如一個(gè)最簡(jiǎn)單的客戶端：

#/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)

試著運(yùn)行一下，可以發(fā)現(xiàn)它成功的打印出了推理結(jié)果：

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

寫在最后

在上述案例中，用戶因時(shí)間緊迫沒法做充分調(diào)研，抱著試試看的想法選擇了 OneFlow，幸運(yùn)的是，終于在極限壓縮的項(xiàng)目周期里順利完成了任務(wù)。

基于 OneFlow v0.7.0 ，用戶輕松地將之前開發(fā)的 PyTorch 的業(yè)務(wù)模型代碼一鍵遷移成 OneFlow 的模型代碼，再經(jīng)過簡(jiǎn)單加工就轉(zhuǎn)成 OneFlow 的靜態(tài)圖 nn.Graph 模式，并利用 nn.Graph 豐富、高效、簡(jiǎn)潔的優(yōu)化開關(guān)來快速大幅提升模型的訓(xùn)練速度，利用完善的周邊工具鏈如 OneFlow-Serving 方便的進(jìn)行線上部署。值得一提的是，用戶還可以使用 OneFlow-ONNX 工具將 OneFlow 高效訓(xùn)練好的模型轉(zhuǎn)成 ONNX 格式導(dǎo)入到其他框架中使用。

本文只介紹了借助和 PyTorch 的兼容性 OneFlow 幫助用戶實(shí)現(xiàn)模型加速和部署的例子。OneFlow 原來的殺手锏功能“大規(guī)模分布式”還沒有體現(xiàn)出來，未來，我們將進(jìn)一步介紹 OneFlow 如何幫助習(xí)慣 PyTorch 的用戶便捷地實(shí)現(xiàn)大規(guī)模預(yù)訓(xùn)練 Transformer 模型和搜索推薦廣告領(lǐng)域需要的大規(guī)模 embedding 模型。

OneFlow項(xiàng)目地址：https://github.com/Oneflow-Inc/oneflow/OneFlow用戶文檔：https://docs.oneflow.org/master/index.html