兼容PyTorch,25倍性能加速,国产框架OneFlow「超速」了(1)

简介: 兼容PyTorch,25倍性能加速,国产框架OneFlow「超速」了

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


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.Graphnn.Module 对象,在 nn.Graph 里执行时,就会采用静态图模式执行,如此动态图下的计算逻辑就可以被静态图直接复用,这样就实现了动静执行的切换。特殊一点的是,Optimizer 也可以添加进入静态图,这样前向、后向、模型更新可以被加入一个完整的静态图做联合优化。

下面的步骤把动态执行的 ResNet101Module 变成静态执行,使用方式和 nn.Module类似,只需要声明、实例化、调用三个基本步骤。

1)声明一个静态图:主要包括两部分,先在初始化函数中添加要静态化的 nn.ModuleOptimizer;然后在 build 函数中构图。



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 上也可以做到。

相关实践学习
基于阿里云DeepGPU实例,用AI画唯美国风少女
本实验基于阿里云DeepGPU实例,使用aiacctorch加速stable-diffusion-webui,用AI画唯美国风少女,可提升性能至高至原性能的2.6倍。
相关文章
|
15天前
|
存储 人工智能 PyTorch
基于PyTorch/XLA的高效分布式训练框架
基于PyTorch/XLA的高效分布式训练框架
25 2
|
8天前
|
机器学习/深度学习 PyTorch 算法框架/工具
PyTorch与迁移学习:利用预训练模型提升性能
【4月更文挑战第18天】PyTorch支持迁移学习,助力提升深度学习性能。预训练模型(如ResNet、VGG)在大规模数据集(如ImageNet)训练后,可在新任务中加速训练,提高准确率。通过选择模型、加载预训练权重、修改结构和微调,可适应不同任务需求。迁移学习节省资源,但也需考虑源任务与目标任务的相似度及超参数选择。实践案例显示,预训练模型能有效提升小数据集上的图像分类任务性能。未来,迁移学习将继续在深度学习领域发挥重要作用。
|
2月前
|
机器学习/深度学习 人工智能 PyTorch
Python 与人工智能的完美结合——解析 PyTorch 框架
【2月更文挑战第4天】本文将探讨 Python 在人工智能领域中的应用,以及介绍 PyTorch 框架。PyTorch 是一个基于 Python 的开源机器学习库,其强大的自动微分功能和易于使用的接口使其成为深度学习领域的热门选择。本文将从 PyTorch 的发展历程、工作原理以及示例代码等方面进行详细分析和解释。
|
5月前
|
机器学习/深度学习 PyTorch 算法框架/工具
深度学习Pytorch框架Tensor张量
深度学习Pytorch框架Tensor张量
45 0
|
6月前
|
机器学习/深度学习 数据可视化 PyTorch
PyTorch 与 TensorFlow:机器学习框架之战
PyTorch 与 TensorFlow:机器学习框架之战
160 0
|
8月前
|
机器学习/深度学习 PyTorch 算法框架/工具
机器学习框架PyTorch详解和案列分析
PyTorch 是一个基于 Python 的机器学习框架,由 Facebook 于 2016 年发布。它提供了一组灵活且高效的工具,可用于构建和训练各种深度学习模型。PyTorch 的核心组件是张量,它是一个多维数组,可以用于存储和处理数据。PyTorch 的张量与 NumPy 的数组类似,但也提供了 GPU 加速和自动微分等功能。PyTorch 使用动态计算图,这意味着在运行时可以修改计算图,从而允许更灵活的模型构建和调试。这与 TensorFlow 等框架的静态计算图不同。PyTorch 支持自动微分,可以方便地计算张量的梯度。这为构建和训练深度学习模型提供了便利。PyTorch 提供了构
246 0
|
8月前
|
数据采集 自然语言处理 PyTorch
全套解决方案:基于pytorch、transformers的中文NLP训练框架,支持大模型训练和文本生成,快速上手,海量训练数据!
全套解决方案:基于pytorch、transformers的中文NLP训练框架,支持大模型训练和文本生成,快速上手,海量训练数据!
全套解决方案:基于pytorch、transformers的中文NLP训练框架,支持大模型训练和文本生成,快速上手,海量训练数据!
|
9月前
|
机器学习/深度学习 人工智能 架构师
AI框架跟计算图什么关系?PyTorch如何表达计算图?
目前主流的深度学习框架都选择使用计算图来抽象神经网络计算表达,通过通用的数据结构(张量)来理解、表达和执行神经网络模型,通过计算图可以把 AI 系统化的问题形象地表示出来。 本节将会以AI概念落地的时候,遇到的一些问题与挑战,因此引出了计算图的概念来对神经网络模型进行统一抽象。接着展开什么是计算,计算图的基本构成来深入了解诶计算图。最后简单地学习PyTorch如何表达计算图。
112 0
|
10月前
|
存储 机器学习/深度学习 缓存
阿里云PAIx达摩院GraphScope开源基于PyTorch的GPU加速分布式GNN框架
阿里云机器学习平台 PAI 团队和达摩院 GraphScope 团队联合推出了面向 PyTorch 的 GPU 加速分布式 GNN 框架 GraphLearn-for-PyTorch(GLT) 。
阿里云PAIx达摩院GraphScope开源基于PyTorch的GPU加速分布式GNN框架
|
11月前
|
机器学习/深度学习 PyTorch 测试技术
RepVGG溯源 | RepVGG作者前期作品,ACNet零算力白嫖性能(附Pytorch代码详解)(二)
RepVGG溯源 | RepVGG作者前期作品,ACNet零算力白嫖性能(附Pytorch代码详解)(二)
74 0