OneFlow v0.6.0 正式釋出

語言: CN / TW / HK

​今天是 OneFlow 開源的 529 天,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

undie2. 實現了與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/