兼容PyTorch，25倍性能加速，国产框架OneFlow「超速」了

计算机视觉研究院

发布于 2022-05-17 16:42:09

5330

发布于 2022-05-17 16:42:09

文章被收录于专栏：计算机视觉战队

关注并星标

从此不迷路

计算机视觉研究院

公众号ID｜ComputerVisionGzq

学习群｜扫码在主页获取加入方式

计算机视觉研究院专栏

作者：Edison_G

要想炼丹爽得飞起，就要选一个顺手的炉子。

转自《机器之心》

要想炼丹爽得飞起，就要选择一个顺手的炉子。作为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.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：

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

© THE END

转载请联系本公众号获得授权

计算机视觉研究院学习群等你加入！

计算机视觉研究院主要涉及深度学习领域，主要致力于人脸检测、人脸识别，多目标检测、目标跟踪、图像分割等研究方向。研究院接下来会不断分享最新的论文算法新框架，我们这次改革不同点就是，我们要着重”研究“。之后我们会针对相应领域分享实践过程，让大家真正体会摆脱理论的真实场景，培养爱动手编程爱动脑思考的习惯！