OneFlow v0.6.0 正式发布

今天是 OneFlow 开源的 528 天，OneFlow v0.6.0 正式发布。 点击“ 阅读原文 ”，欢迎下载体验最新版本。本次版本更新包括框架、模型和 OneFlow-ONNX 三大部分，主要有：

• 性能提升，包括静态图、动态图、算子性能、显存占用等方面

• 新增大量常用算子

• 完善静态图和 ConsistentTensor 功能

• 支持 OneFlow 作为 Nvidia Triton 的后端提供 Serving 功能

• 实现丰富的视觉预训练模型，与 torchvision、timm 对齐

• 实现更加完善的 OneFlow-ONNX 转换功能

以下为版本更新详情。

框架优化

1. 深度优化 nn.Graph 的性能

• 与 v0.5.0 相比， v0.6.0 的 nn.Graph 在 ResNet AMP  和 WDL 等模型上的训练速度提升了 10%

  • 新版的动静转换功能的性能还有可优化的空间，近期着重优化了 nn.Graph 在高频率迭代训练场景下的性能

  • 重新设计实现了 nn.Graph 的调度指令， 重构了 Actor Graph 与 Eager VM 的交互逻辑，使得 Graph 的 runtime 执行与 Python input/output Tensor 尽可能异步流水并行

2. 深度优化 Eager 性能

• 与 v0.5.0 相比，v0.6.0 OneFlow Eager 在小 batch 场景下的训练速度大幅提升

  • 深度优化虚拟机的调度逻辑

  • 优化 get/set item

  • 优化 tensor.numel()

  • 优化 oneflow.Size()

3. 深度优化算子性能

• 着重优化了一些影响新模型性能瓶颈的算子，使得相关的模型训练速度有明显提升

  • 新增 fused dropout系列算子

  • 新增 CPU版 group deconv 并优化性能

  • 给以下算子新增 inplace 版本实现：mul, hard_sigmoid, sin

  • 对linalg.vector_norm在ord=2.0时进行了性能优化，相比之前提升了4倍

  • 深度优化LayerNorm op，性能大幅度领先PyTorch和Apex实现

  • 实现op的数据类型自动提升

4. 深度优化 Eager 显存占用

• 优化了某些算子在网络训练中对显存占用，使得相同计算设备可以跑更大的模型或数据

  • 优化 broadcast binary 一族算子的后向显存占用

  • 优化 Slice 算子的后向显存占用

  • 优化 LayerNorm 的显存占用  

5. 给静态图 nn.Graph 新增众多实用功能

• 静态图抽象 nn.Graph 增加了许多新功能，涉及静态图的效率、调试、完备性以及在更多场景下的易用性等方面：

  • 为了辅助静态图的调试，我们新增了：

  • debug 模式支持 graph.debug(1) 打印更多构图信息

  • 提供环境变量：ONEFLOW_DEBUG_PASS 来显示编译期 图优化前后计算图的变化

  • 给 Nsight Profile 增加用户可读的线程命名信息，方便定位和检索目标关键线程位置

  • 丰富了大量静态图的测试用例：增加伴随Eager测试的自动nn.Graph测试

  • 为了支持使用 nn.Graph 做模型的部署（Serving），提供了 graph.save() 和 load() 接口

  • 为了在使用 TensorCore 的 GPU 上做 AMP 的加速，提供了环境变量：ONEFLOW_ENABLE_NHWC 用于表示 CNN 相关算子进行 channels last 计算

  • 使得 nn.Graph 支持更多的使用场景：

  • 支持 稀疏更新 Optimizer，用于 WDL 场景下的参数稀疏更新

  • 支持在nn.Graph下使用Sequential, ModuleList, ModuleDict, ParameterList, ParameterDict这些nn.Module Container

  • 支持在nn.Graph的init函数中创建Optimizer

  • 支持nn.Graph下多个参数共用同一个Tensor

  • 支持实际的进程数大于GPU设备数的使用场景

  • nn.Graph下Consistent的SBP推理时考虑Inplace，支持更多Inplace执行

6. 新增了大量算子

• 新增算子：cumsum, meshgrid, linspace, diagonal, movedim, roialign, nms, arccos, roll

• 新增算子：masked_fill, floordiv, glu, pool1d, pool2d, pool3d

• 新增unfold和fold op

• 实现op的数据类型自动提升

• 实现expand和repeat op

• 目前 torchvision 库的模型可以通过import oneflow as torch实现一键切换

7. 支持用户自定义 autograd.Function

• 用户可以像 Torch 一样自定义 autograd.Function

8. 提供基础的 Serving 功能

• 支持 OneFlow 作为 Triton 的 backend 提供模型的 Serving 功能

9. 新增 Tensor（ConsistentTensor） 的部分功能

• 支持 Tensor 使用 2-D SBP 来表示任意的混合并行方式（如一个 Linear 运算在设备矩阵的行方向上数据并行，在列方向上模型并行）

• 支持 Tensor 从任意的 1-D SBP 到 2-D SBP 的转换（网络由 1-D 并行 和 2-D 并行混合组成）

• 支持从 numpy 构造 ConsistentTensor

• 新增 oneflow.from_numpy()

• 新增 oneflow.numel()

• 新增 tensor.expand_as()  ###  

模型实现

发布 flowvison  0.0.54 

（链接：https://github.com/Oneflow-Inc/vision）

1. 实现了丰富的视觉预训练模型

图像分类

• CNN系列: ResNet, DenseNet, VGG, ResNext, EfficientNet等

• Vision Transformer系列: ViT, PVT, Swin-Transformer等

• Vision MLP系列:Mlp-Mixer, Res-MLP, g-MLP等

目标检测

• SSD, SSDLite

• Faster R-CNN

• RetinaNet

图像分割

• FCN

• DeepLabV3

风格迁移

• StyleNet: 支持风格sketch, candy, mosaic, rain_princess, undie

2. 实现了与torchvision对齐的数据增强操作

包括 CenterCrop , ColorJitter 等与torchvision对齐的数据增强操作，在大多数场景下可以 import flowvision as torchvision 直接替换

3. 对齐了timm中的高级的数据增强实现

flowvision.data中所实现的高级数据增强操作

• Mixup

• CutMix

• Random-Erasing

• AutoAugment

• RandAugment

• AugMix

4. 单独抽离出Layers模块，提供搭建模型时即插即用的Block

flowvision.layers.attention模块

• 实现了Non-Local, SELayer, CBAM, BAM, ECA等即插即用的attention模块

flowvision.layers.blocks模块

• 提供了PatchEmb, Pooler, ConvBnAct等在搭建模型时可能用到的模块

flowvision.layers.regularization模块

• 提供了drop-path, drop-block, stochastic depth等正则化模块，用来提升模型泛化能力  此外还有activation, weight_init等单独的文件，用来提供激活函数，初始化方法等组件    

OneFlow-ONNX转换

更新 OneFlow 转 ONNX 模型格式的 工具 包

• 支持CPU和GPU模式的OneFlow模型转onnx模型

• 新增算子和模型测试样例，对齐OneFlowVision库中的全部分类模型

• 修复PReLU转换时出现的onnx-runtime相关的bug

• 兼容1.9.0版本以上的onnx-runtime库

• 发布0.5.4版本oneflow-onnx包，pip install oneflow-onnx即可体验

欢迎下载体验OneFlow新一代开源深度学习框架：https://github.com/Oneflow-Inc/oneflow/