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 的进一步的使用可以参考:
- nn.Graph 的使用教程,https://docs.oneflow.org/en/master/basics/08_nn_graph.html
- nn.Graph 的 API 文档,https://oneflow.readthedocs.io/en/master/graph.html
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 模型。