OneFlow｜推荐一款兼容 PyTorch 的国产 AI 框架

在人工智能和深度学习的不断发展进程中，深度学习框架已经成为行业研发基础，承载着越来越多的研发者基于框架构建、训练和部署相关应用。基于深度学习框架，研发人员可以依据研究或业务目标更快速地开发相关应用，行业的研发效率和水平也因此得到了显著提升。在深度学习框架中，PyTorch 和 TensorFlow 推出时间比较早，在科研和研发场景渗透率都比较高，PyTorch 是目前受到用户喜欢的框架之一，它的 API 友好，Eager 模式让模型搭建和调试过程变得更加容易，但是静态图编译和部署体验还不能完全令人满意；另外一款用户常用的框架 TensorFlow ，其静态编译和部署功能很完备，但调试体验相对复杂。

面对现有框架的优劣势，OneFlow 在2020 年 7 月在 GitHub 上开源时，差异化地定位于框架对分布式和高性能的支持，为了更好地满足研发场景的需求，团队用一年多时间自研了动态图引擎，旨在兼容并包 PyTorch 与 TensorFlow 的优势。OneFlow 0.7 版本已支持与 PyTorch 一致的 Eager 体验，实现了同时支持动态图和静态图。不仅如此，OneFlow 编程 API 与 PyTorch 兼容，常见深度学习模型只需修改一行 import oneflow as torch 就可以把 PyTorch 写的模型在 OneFlow 上跑起来。目前，矩池云 Matpool 已经支持 OneFlow 框架，只需要在训练时选择框架，就可以体验到这款优秀的国产框架。

OneFlow 视觉模型库flowvision

OneFlow 视觉模型库 flowvision 已经支持计算机视觉领域图像分类、分割和检测等方向的经典 SOTA 模型 （见下表），这些模型都可以通过 import torch as flow 或 import oneflow as torch 实现自由切换。

OneFlow 和 PyTorch 的兼容，意味着用户可以像使用 PyTorch 一样来使用 OneFlow ，如果对模型效果比较满意，就可以继续使用 OneFlow 扩展到大规模分布式或使用静态图部署模型。

迁移调优案例

项目案例背景

一家头部通信公司将上线一款基于深度学习的图像识别应用，项目基于 PyTorch 建立了业务模型，且已经跑通了正常训练流程，但是训练/推理速度都很慢，在本次案例中，项目方研发人员将模型脚本迁移到 OneFlow ，在 OneFlow 团队的支持下，进行了大幅度的训练/推理性能优化、部署上线。

研发人员可以借助与 PyTorch 兼容的 OneFlow 实现模型加速和部署，** OneFlow 发布的官方案例展示了通过多方位调优，达到整体性能提升的一些方法，其中涉及到的调优思路，既有共性的调优方法，也有在 OneFlow 框架下所特有调优方式 ** 。

项目案例迁移调优过程

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

OneFlow 最新发布的 0.7.0 版本对 PyTorch 接口的兼容性有了进一步的完善。OneFlow 对已经支持的算子都能保证和 PyTorch 的接口在语义和结果上一致。在本案例中，业务模型的主干网络是 resnet101，在迁移过程中，可参考 官方文档 迁移 ，只需要模型文件中与 torch 相关的 import 修改为 import oneflow as torch，就完成了模型代码的迁移工作。

在模型脚本迁移完毕之后，还需要验证模型迁移的准确性，确定精度是不是对齐了。

1.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)

1.2 训练流程的验证

在验证完推理精度后是验证训练流程，在对齐训练超参数之后，使用 OneFlow 训练模型的 loss 曲线和 PyTorch 的 loss 曲线也高度一致，在小数据集上的精度高度一致。

2 使用 OneFlow 的 nn.Graph 加速模型训练与推理性能

在验证完算法正确性后，需要考虑的是如何进行加速执行。

2.1 动态转静态

OneFlow 有 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类似，只需要声明、实例化、调用三个基本步骤。

• 声明一个静态图：主要包括两部分，先在初始化函数中添加要静态化的 nn.Module和 Optimizer；然后在 build函数中构图。