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

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
相关文章
|
8月前
|
存储 人工智能 PyTorch
基于PyTorch/XLA的高效分布式训练框架
基于PyTorch/XLA的高效分布式训练框架
335 2
|
27天前
|
人工智能 安全 PyTorch
SPDL:Meta AI 推出的开源高性能AI模型数据加载解决方案,兼容主流 AI 框架 PyTorch
SPDL是Meta AI推出的开源高性能AI模型数据加载解决方案,基于多线程技术和异步事件循环,提供高吞吐量、低资源占用的数据加载功能,支持分布式系统和主流AI框架PyTorch。
55 10
SPDL:Meta AI 推出的开源高性能AI模型数据加载解决方案,兼容主流 AI 框架 PyTorch
|
2月前
|
机器学习/深度学习 监控 PyTorch
深度学习工程实践:PyTorch Lightning与Ignite框架的技术特性对比分析
在深度学习框架的选择上,PyTorch Lightning和Ignite代表了两种不同的技术路线。本文将从技术实现的角度,深入分析这两个框架在实际应用中的差异,为开发者提供客观的技术参考。
54 7
|
7月前
|
机器学习/深度学习 人工智能 PyTorch
|
3月前
|
机器学习/深度学习 并行计算 PyTorch
提高 PyTorch 性能
提高 PyTorch 是一个非常流行的深度学习框架,它支持动态计算图,非常适合快速原型设计和研究。
55 3
|
4月前
|
机器学习/深度学习 PyTorch 调度
在Pytorch中为不同层设置不同学习率来提升性能,优化深度学习模型
在深度学习中,学习率作为关键超参数对模型收敛速度和性能至关重要。传统方法采用统一学习率,但研究表明为不同层设置差异化学习率能显著提升性能。本文探讨了这一策略的理论基础及PyTorch实现方法,包括模型定义、参数分组、优化器配置及训练流程。通过示例展示了如何为ResNet18设置不同层的学习率,并介绍了渐进式解冻和层适应学习率等高级技巧,帮助研究者更好地优化模型训练。
259 4
在Pytorch中为不同层设置不同学习率来提升性能,优化深度学习模型
|
5月前
|
机器学习/深度学习 并行计算 PyTorch
GPU 加速与 PyTorch:最大化硬件性能提升训练速度
【8月更文第29天】GPU(图形处理单元)因其并行计算能力而成为深度学习领域的重要组成部分。本文将介绍如何利用PyTorch来高效地利用GPU进行深度学习模型的训练,从而最大化训练速度。我们将讨论如何配置环境、选择合适的硬件、编写高效的代码以及利用高级特性来提高性能。
936 1
|
5月前
|
机器学习/深度学习 人工智能 PyTorch
AI智能体研发之路-模型篇(五):pytorch vs tensorflow框架DNN网络结构源码级对比
AI智能体研发之路-模型篇(五):pytorch vs tensorflow框架DNN网络结构源码级对比
91 1
|
5月前
|
UED 开发者
哇塞!Uno Platform 数据绑定超全技巧大揭秘!从基础绑定到高级转换,优化性能让你的开发如虎添翼
【8月更文挑战第31天】在开发过程中,数据绑定是连接数据模型与用户界面的关键环节,可实现数据自动更新。Uno Platform 提供了简洁高效的数据绑定方式,使属性变化时 UI 自动同步更新。通过示例展示了基本绑定方法及使用 `Converter` 转换数据的高级技巧,如将年龄转换为格式化字符串。此外,还可利用 `BindingMode.OneTime` 提升性能。掌握这些技巧能显著提高开发效率并优化用户体验。
72 0
|
5月前
|
机器学习/深度学习 PyTorch TensorFlow
深度学习框架之争:全面解析TensorFlow与PyTorch在功能、易用性和适用场景上的比较,帮助你选择最适合项目的框架
【8月更文挑战第31天】在深度学习领域,选择合适的框架至关重要。本文通过开发图像识别系统的案例,对比了TensorFlow和PyTorch两大主流框架。TensorFlow由Google开发,功能强大,支持多种设备,适合大型项目和工业部署;PyTorch则由Facebook推出,强调灵活性和速度,尤其适用于研究和快速原型开发。通过具体示例代码展示各自特点,并分析其适用场景,帮助读者根据项目需求和个人偏好做出明智选择。
129 0