关于人工智能:兼容PyTorch25倍性能加速OneFlow超速

起源|机器之心

要想炼丹爽得飞起，就要抉择一个棘手的炉子。作为 AI 工程师日常必不可缺的「炼丹炉」，「PyTorch 还是 TensorFlow？」已成为知乎、Reddit 等炼丹师出没之地每年都会探讨的热门话题。

业界流传一种说法：PyTorch 适宜学术界，TensorFlow 适宜工业界。毕竟，PyTorch 是用户最喜爱的框架，API 十分敌对，Eager 模式让模型搭建和调试过程变得更加容易，不过，它的动态图编译和部署体验还不令人满意。TensorFlow 恰恰相反，动态编译和部署性能很齐备，不过其调试体验让人欲哭无泪。

那么问题来了： 鱼和熊掌真的不可兼得吗？ 未必，来自北京的一流科技团队推出的开源深度学习框架 OneFlow 曾经做到了。

等等，OneFlow 始终主打分布式和高性能，易用性也能和 PyTorch 一样吗？据说过 OneFlow 的人肯定会收回这样的疑难。

没错，从 2016 年底立项之日起，OneFlow 就是为大规模分布式而生，特色之一就是动态图机制，2020 年 7 月在 GitHub 上开源时还不反对动态图。不过，OneFlow 团队用一年多工夫自研了动态图引擎，OneFlow v0.7.0 版本已反对和 PyTorch 截然不同的 Eager 体验，也就是说，OneFlow 实现了同时反对动态图和动态图。不仅如此，OneFlow 编程 API 齐全和 PyTorch 兼容，常见深度学习模型只需批改一行 import oneflow as torch 就能够把 PyTorch 写的模型在 OneFlow 上跑起来。

无妨先到 OneFlow 视觉模型库 flowvision 看一看：https://github.com/Oneflow-In…，这个模型库曾经反对计算机视觉畛域图像分类、宰割和检测等方向的经典 SOTA 模型（见下表），这些模型都能够通过 import torch as oneflow 或 import oneflow as torch 实现自在切换。

OneFlow 和 PyTorch 兼容之后，用户能够像应用 PyTorch 一样来应用 OneFlow，对模型成果比较满意之后，能够持续应用 OneFlow 扩大到大规模分布式或应用动态图部署模型。听下来是不是 too good to be true？

在上面的案例中，一家头部通信公司基于 PyTorch 的业务模型疾速不便地迁徙成 OneFlow 的模型，并进行大幅度的训练 / 推理性能优化、部署上线，短短几天工夫就让业务得以按时上线部署，且各项性能指标均大幅超出预期！

他们到底是如何做到的？先从我的项目背景说起。

1

为什么抉择 OneFlow？

因业务倒退需要，这家通信公司近期将上线一款基于深度学习的图像识别利用，该项目标业务需要有如下五个特点：

• 数据量大：数据库中有过亿级别的图片
• 模型简略：比拟惯例的分类模型
• 400 多张显卡，短期内无奈扩容
• 对于训练 / 推理的吞吐有硬性指标
• 上线工夫紧迫

用户基于市面上最风行的深度学习框架 PyTorch 搭建了业务模型，且跑通了失常训练流程，然而训练 / 推理都很慢，远远达不到指标（离上线 QPS 有 20 倍的差距），随着交付日期邻近，整个团队深陷焦虑。

用户尝试了各种计划（基于已有实现进行优化）都杯水车薪，于是调研了其余深度学习框架，如 TensorFlow、OneFlow 等，发现 OneFlow (https://github.com/OneFlow-In…) 是减速 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/mast…）以及 OneFlow 研发团队的大力支持下，用户发展了以下工作：

将已有 PyTorch 的我的项目代码齐全迁徙到 OneFlow

将我的项目代码由动态图模式（Eager Mode）革新为动态图模式（Graph Mode）

开启 OneFlow Graph 模式下的各种优化选项并训练模型

用 Serving 模块部署模型上线

2

迁徙调优过程

1. 一键迁徙 PyTorch 模型转 OneFlow 模型：只需 import oneflow as torch 就够了

OneFlow 最新公布的 0.7.0 版本对 PyTorch 接口的兼容性有了进一步的欠缺。OneFlow 对曾经反对的算子都能保障和 PyTorch 的接口在语义和后果上统一。于是用户就尝试了一下迁徙模型脚本到 OneFlow。因为业务模型的骨干网络是 resnet101，在迁徙过程中，用户参考了官网文档（https://docs.oneflow.org/mast…）来迁徙，发现只须要模型文件中与 torch 相干的 import 批改为 import oneflow as torch，就实现了模型代码的迁徙工作。

在模型脚本迁徙结束之后，还须要验证模型迁徙的正确性，看看精度是不是对齐了。

1）用户首先做了推理精度的验证，就是间接加载 PyTorch 训练好的模型而后验证推理精度，因为 OneFlow 对齐了 PyTorch 的接口，所以加载 PyTorch 的模型也十分不便，只需数行代码即可实现：

import torchvision.models as models_torch
import flowvision.models as models_flow

resnet101_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-In…），理解到单卡的执行效率曾经做得很高，照猫画虎，这些优化技巧都能够用在 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/deepl…）。不同的 input 和 filter 大小在不同的算法下有不同的性能体现，为了抉择最佳算法，在调用 cudnn convolution 算子接口前，须要先调用 cudnn convolution searching algorithm 的接口。cudnn 提供了 2 种搜寻模式：启发式搜（https://docs.nvidia.com/deepl…）和试运行搜寻（cudnnFindConvolutionForwardAlgorithm）（https://docs.nvidia.com/deepl…）。

启发式搜寻是通过一种「查表」的形式来搜查最佳算法，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/m…
• nn.Graph 的 API 文档：https://oneflow.readthedocs.i…

3. 应用 OneFlow-Serving，轻松将训练好的模型部署上线

当用户实现训练，失去最终的模型之后，接下来的一步就是模型部署。不同于模型训练时须要进行权重更新，部署时的权重固定不变，所以能够进行更激进的速度优化，例如 int8 量化、更宽泛的 kernel fusion、constant folding 等等。

用户参考 OneFlow v0.7.0 提供了官网的 Serving 模块（https://github.com/Oneflow-In…），它是一个 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 np
import tritonclient.http as httpclient
from 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 / 255
image = np.expand_dims(image, axis=0)
# Transpose NHWC to NCHW
image = 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

3

写在最初

在上述案例中，用户因工夫紧迫没法做充沛调研，抱着试试看的想法抉择了 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-In…
• OneFlow 用户文档：https://docs.oneflow.org/mast…

欢送下载体验 OneFlow v0.7.0 最新版本：

https://github.com/Oneflow-In…