要想炼丹爽得飞起,就要选择一个顺手的炉子。作为 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.Graph 的 nn.Module 对象,在 nn.Graph 里执行时,就会采用静态图模式执行,如此动态图下的计算逻辑就可以被静态图直接复用,这样就实现了动静执行的切换。特殊一点的是,Optimizer 也可以添加进入静态图,这样前向、后向、模型更新可以被加入一个完整的静态图做联合优化。
下面的步骤把动态执行的 ResNet101Module 变成静态执行,使用方式和 nn.Module类似,只需要声明、实例化、调用三个基本步骤。
1)声明一个静态图:主要包括两部分,先在初始化函数中添加要静态化的 nn.Module 和 Optimizer;然后在 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 上也可以做到。