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

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

2.4 静态编译优化

做到以上优化后,GPU 利用率已经能比较稳定的保持在 90%,一般来说,已经没有太大优化空间了。但是,OneFlow nn.Graph 下还有一些自动的编译优化技术可以尝试。

比如利用自动混合精度做低精度计算、利用算子融合来减少访存开销等,这里最终带来了 64% 的加速,速度到了原来最好性能的 1.56 倍。

此前示例中提到的 _config_graph 函数就是在配置这些优化选项,具体如下:



def _config_graph(graph):    if args.fp16:        # 打开 nn.Graph 的自动混合精度执行        graph.config.enable_amp(True)     if args.conv_try_run:        # 打开 nn.Graph 的卷积的试跑优化        graph.config.enable_cudnn_conv_heuristic_search_algo(False)     if args.fuse_add_to_output:        # 打开 nn.Graph 的add算子的融合        graph.config.allow_fuse_add_to_output(True)     if args.fuse_pad_to_conv:        # 打开 nn.Graph 的pad算子的融合        graph.config.allow_fuse_pad_to_conv(True)


对于 ResNet101,batch_size 设置为 16,在 nn.Graph 无优化选项打开的基础上:

  • 打开混合精度,测试得到了 36% 的加速

 自动混合精度训练,自动将网络中的合适的算子由 FP32 单精度计算转换成 FP16 半精度浮点进行计算,不仅可以减少 GPU 显存占用,而且可以提升整体性能,在支持 Tensor Core 的 GPU 设备上还会使用 Tensor Core 进一步加速训练。

  • 再打开卷积试跑优化,测试得到了 7% 的加速,总加速为 43%


cudnn 的 convolution 算子包含多种算法,例如前向的算法(https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnConvolutionFwdAlgo_t)。不同的 input 和 filter 大小在不同的算法下有不同的性能表现,为了选择最佳算法,在调用 cudnn convolution 算子接口前,需要先调用 cudnn convolution searching algorithm 的接口。cudnn 提供了2种搜索模式:启发式搜索(https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnGetConvolutionForwardAlgorithm_v7)和试运行搜索(cudnnFindConvolutionForwardAlgorithm)(https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnFindConvolutionForwardAlgorithm)。

启发式搜索是通过一种「查表」的方式来搜寻最佳算法,cudnn 对不同的参数配置对应的最佳算法进行了预先定义,然后每次搜索时进行匹配得到结果。试运行搜索会传入实际的张量进行多次试运行,然后返回运行结果。搜索算法返回的结果都是不同算法的元信息及其所需耗时。

启发式搜索在搜索阶段不需额外分配内存,且能更快得到结果;而试运行搜索能得到更为全面和精确的结果,也即通常能更精确地找到最佳算法。启发式搜索在常见情形下可以得到与试运行搜索一致的结果,但在一些特殊参数配置下无法得到最佳结果。OneFlow 中默认启动了启发式搜索,但可通过 graph.config.enable_cudnn_conv_heuristic_search_algo(False) 接口关闭,关闭后使用的就是试运行搜索。

  • 再打开 pad 和 conv 算子融合,测试得到了 19% 的加速,总加速为 62%


在 CNN 网络 Backbone 中有很多 convolution + pad 的组合,convolution 算子自身支持 pad 操作,自动将 pad 算子 fuse 到 convolution 算子上,可以省掉 pad 算子的开销,提升网络整体性能。

  • 再打开 add 的算子的融合,测试得到了 2% 的加速,总加速为 64%


自动将网络中常见的访存密集型算子 Elementwise add 算子和上游的算子 fuse 起来,可以减少带宽使用,从而提升性能。对于 Elementwise add 算子来说,将其 fuse 到上一个算子,可以减少一次数据读写,有约 2/3 的性能提升。

另外 nn.Graph 可以很方便地支持使用 TensorRT 。本优化对象没有更新模型的需求,所以也适合使用 TensorRT 做加速。在 nn.Graph 无优化选项基础上, batch_size 设置为 16,新增自动混合精度、NHWC、使用 TensorRT 后端,可以提速 48%。

在这个模型里,只使用 TensorRT 后端比只使用 OneFlow 的静态图优化还差一点,可能的原因是, TensorRT 下的一些优化在 nn.Graph 里已经做了,所以没有带来额外收益。不过其实验起来还比较方便,编译一下带 TensorRT 的 OneFlow,再在 nn.Graph 下打开开关就可以,列出来作为参考:



def _config_graph(graph):    if args.tensorrt:        # 使用 TensorRT 后端执行        graph.config.enable_tensorrt(True)


2.5 加速优化总结

以上记录了加速的主要过程,动态转静态加速约 1.25 倍、算法逻辑约减加速约 2.33 倍、提高并行度加速约 4.6 倍、静态编译优化加速约 1.6 倍,累积加速约 21 倍。中间有些小的优化点没有完全记录,实际累积的加速效果达到了 25 倍以上,超过了项目部署的 20 倍加速需求。

nn.Graph 的进一步的使用可以参考:


3. 使用 OneFlow-Serving,轻松将训练好的模型部署上线

当用户完成训练,得到最终的模型之后,接下来的一步就是模型部署。不同于模型训练时需要进行权重更新,部署时的权重固定不变,所以可以进行更激进的速度优化,例如 int8 量化、更广泛的 kernel fusion、constant folding 等等。

用户参考 OneFlow v0.7.0 提供了官方的 Serving 模块(https://github.com/Oneflow-Inc/serving),它是一个 NVIDIA Triton 的后端,集成了 OneFlow 内置的 XRT 模块,并提供了开箱即用的用户接口。只需使用下述方法就将训练好的 OneFlow 模型快速高效的部署起来:

为了将模型用于推理,在使用 nn.Graph 训练完成之后,需要构造一个只包含前向的 ResNet101InferenceGraph


class ResNet101InferenceGraph(oneflow.nn.Graph):    def __init__(self):        super().__init__()                self.model = resnet101_graph.model        def build(self, input):        return self.model(input) inference_graph = ResNet101InferenceGraph()


并以一个样例输入运行 inference_graph,触发 inference_graph 的计算图构建:


unused_output = inference_graph(flow.zeros(1, 3, 224, 224))


接下来就可以运行 flow.save inference_graph 的计算图结构以及权重均保存在 "model" 文件夹下,以供部署使用:


flow.save(inference_graph, "model")


然后只需要运行




docker run --rm --runtime=nvidia --network=host -v$(pwd)/model:/models/resnet101/1 \  oneflowinc/oneflow-serving:nightly


由此可以启动一个部署着 ResNet101 模型的 Docker 容器。这里的 -v 很重要,它表示将当前目录下的 model 文件夹映射到容器内的 "/models/resnet101/1" 目录,其中 /models 是 Triton 读取模型的默认目录,Triton 会以该目录下的一级目录名("resnet101")作为模型名称,二级目录名("1")作为模型版本。

如果将启动命令调整为



docker run --rm --runtime=nvidia --network=host -v$(pwd)/model:/models/resnet101/1 \  oneflowinc/oneflow-serving:nightly oneflow-serving --model-store /models --enable-tensorrt resnet101


模型就会通过 OneFlow 的 XRT 模块自动使用 TensorRT 进行推理,此外 OneFlow Serving 还支持类似的 “--enable-openvino”。

启动 Docker 容器后,运行下面的命令,就可以查看服务状态:


curl -v localhost:8000/v2/health/ready


返回值为 HTTP/1.1 200 OK,表示服务正在正常工作。

接下来就可以使用 Triton 的 C++ 或 Python SDK 实现向服务端发送请求并获取结果的逻辑了,例如一个最简单的客户端:




#/usr/bin/env python3 import numpy as npimport tritonclient.http as httpclientfrom PIL import Image  triton_client = httpclient.InferenceServerClient(url='127.0.0.1:8000') image = Image.open("image.jpg")image = image.resize((224, 224))image = np.asarray(image)image = image / 255image = np.expand_dims(image, axis=0)# Transpose NHWC to NCHWimage = np.transpose(image, axes=[0, 3, 1, 2])image = image.astype(np.float32) input = httpclient.InferInput('INPUT_0', image.shape, "FP32")input.set_data_from_numpy(image, binary_data=True)output_placeholder = httpclient.InferRequestedOutput('OUTPUT_0', binary_data=True, class_count=1)output = triton_client.infer("resnet101", inputs=[input], outputs=[output_placeholder]).as_numpy('OUTPUT_0')print(output)


试着运行一下,可以发现它成功的打印出了推理结果:



$ python3 triton_client.py[b'3.630257:499']    # class id 为 499,值为 3.630257


写在最后

在上述案例中,用户因时间紧迫没法做充分调研,抱着试试看的想法选择了 OneFlow,幸运的是,终于在极限压缩的项目周期里顺利完成了任务。

基于 OneFlow v0.7.0 ,用户轻松地将之前开发的 PyTorch 的业务模型代码一键迁移成 OneFlow 的模型代码,再经过简单加工就转成 OneFlow 的静态图 nn.Graph 模式,并利用 nn.Graph 丰富、高效、简洁的优化开关来快速大幅提升模型的训练速度,利用完善的周边工具链如 OneFlow-Serving 方便的进行线上部署。值得一提的是,用户还可以使用 OneFlow-ONNX 工具将 OneFlow 高效训练好的模型转成 ONNX 格式导入到其他框架中使用。

本文只介绍了借助和 PyTorch 的兼容性 OneFlow 帮助用户实现模型加速和部署的例子。OneFlow 原来的杀手锏功能“大规模分布式”还没有体现出来,未来,我们将进一步介绍 OneFlow 如何帮助习惯 PyTorch 的用户便捷地实现大规模预训练 Transformer 模型和搜索推荐广告领域需要的大规模 embedding 模型。

相关文章
|
4月前
|
机器学习/深度学习 PyTorch TensorFlow
66_框架选择:PyTorch vs TensorFlow
在2025年的大语言模型(LLM)开发领域,框架选择已成为项目成功的关键决定因素。随着模型规模的不断扩大和应用场景的日益复杂,选择一个既适合研究探索又能支持高效部署的框架变得尤为重要。PyTorch和TensorFlow作为目前市场上最主流的两大深度学习框架,各自拥有独特的优势和生态系统,也因此成为开发者面临的经典选择难题。
|
5月前
|
机器学习/深度学习 算法 PyTorch
【Pytorch框架搭建神经网络】基于DQN算法、优先级采样的DQN算法、DQN + 人工势场的避障控制研究(Python代码实现)
【Pytorch框架搭建神经网络】基于DQN算法、优先级采样的DQN算法、DQN + 人工势场的避障控制研究(Python代码实现)
156 1
|
5月前
|
机器学习/深度学习 算法 PyTorch
【DQN实现避障控制】使用Pytorch框架搭建神经网络,基于DQN算法、优先级采样的DQN算法、DQN + 人工势场实现避障控制研究(Matlab、Python实现)
【DQN实现避障控制】使用Pytorch框架搭建神经网络,基于DQN算法、优先级采样的DQN算法、DQN + 人工势场实现避障控制研究(Matlab、Python实现)
251 0
|
9月前
|
机器学习/深度学习 PyTorch 编译器
深入解析torch.compile:提升PyTorch模型性能、高效解决常见问题
PyTorch 2.0推出的`torch.compile`功能为深度学习模型带来了显著的性能优化能力。本文从实用角度出发,详细介绍了`torch.compile`的核心技巧与应用场景,涵盖模型复杂度评估、可编译组件分析、系统化调试策略及性能优化高级技巧等内容。通过解决图断裂、重编译频繁等问题,并结合分布式训练和NCCL通信优化,开发者可以有效提升日常开发效率与模型性能。文章为PyTorch用户提供了全面的指导,助力充分挖掘`torch.compile`的潜力。
1055 17
|
机器学习/深度学习 边缘计算 PyTorch
PyTorch团队为TorchAO引入1-8比特量化,提升ARM平台性能
PyTorch团队推出创新技术,在其低精度计算库TorchAO中引入低位运算符支持,实现1至8位精度的嵌入层权重量化及8位动态量化激活的线性运算符。该技术通过模块化设计和高效硬件利用,优化了资源受限环境下的深度学习计算,提升了计算效率并降低了资源消耗。新内核与PyTorch生态系统无缝集成,支持即时执行、编译优化及边缘计算,为开发者提供全方位性能优势。测试结果显示,多层次量化策略显著提升了计算效率,保持了模型精度。这一突破为深度学习框架优化开辟了多个研究方向,推动了人工智能在边缘计算等领域的广泛应用。
476 11
PyTorch团队为TorchAO引入1-8比特量化,提升ARM平台性能
|
人工智能 安全 PyTorch
SPDL:Meta AI 推出的开源高性能AI模型数据加载解决方案,兼容主流 AI 框架 PyTorch
SPDL是Meta AI推出的开源高性能AI模型数据加载解决方案,基于多线程技术和异步事件循环,提供高吞吐量、低资源占用的数据加载功能,支持分布式系统和主流AI框架PyTorch。
700 10
SPDL:Meta AI 推出的开源高性能AI模型数据加载解决方案,兼容主流 AI 框架 PyTorch
|
机器学习/深度学习 并行计算 PyTorch
TorchOptimizer:基于贝叶斯优化的PyTorch Lightning超参数调优框架
TorchOptimizer 是一个基于贝叶斯优化方法的超参数优化框架,专为 PyTorch Lightning 模型设计。它通过高斯过程建模目标函数,实现智能化的超参数组合选择,并利用并行计算加速优化过程。该框架支持自定义约束条件、日志记录和检查点机制,显著提升模型性能,适用于各种规模的深度学习项目。相比传统方法,TorchOptimizer 能更高效地确定最优超参数配置。
658 7
|
12月前
|
机器学习/深度学习 PyTorch TensorFlow
深度学习工具和框架详细指南:PyTorch、TensorFlow、Keras
在深度学习的世界中,PyTorch、TensorFlow和Keras是最受欢迎的工具和框架,它们为研究者和开发者提供了强大且易于使用的接口。在本文中,我们将深入探索这三个框架,涵盖如何用它们实现经典深度学习模型,并通过代码实例详细讲解这些工具的使用方法。
1087 0
|
机器学习/深度学习 PyTorch 调度
在Pytorch中为不同层设置不同学习率来提升性能,优化深度学习模型
在深度学习中,学习率作为关键超参数对模型收敛速度和性能至关重要。传统方法采用统一学习率,但研究表明为不同层设置差异化学习率能显著提升性能。本文探讨了这一策略的理论基础及PyTorch实现方法,包括模型定义、参数分组、优化器配置及训练流程。通过示例展示了如何为ResNet18设置不同层的学习率,并介绍了渐进式解冻和层适应学习率等高级技巧,帮助研究者更好地优化模型训练。
1017 4
在Pytorch中为不同层设置不同学习率来提升性能,优化深度学习模型
|
机器学习/深度学习 监控 PyTorch
深度学习工程实践:PyTorch Lightning与Ignite框架的技术特性对比分析
在深度学习框架的选择上,PyTorch Lightning和Ignite代表了两种不同的技术路线。本文将从技术实现的角度,深入分析这两个框架在实际应用中的差异,为开发者提供客观的技术参考。
435 7

推荐镜像

更多