兼容PyTorch，25倍性能加速，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-Inc/vision，这个模型库已经支持计算机视觉领域图像分类、分割和检测等方向的经典 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-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 研发团队的大力支持下，用户开展了以下工作：

1. 将已有 PyTorch 的项目代码完全迁移到 OneFlow

2. 将项目代码由动态图模式（Eager Mode）改造为静态图模式（Graph Mode）

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

4. 用 Serving 模块部署模型上线

2

迁移调优过程

1. 一键迁