来源|机器之心
要想炼丹爽得飞起,就要选择一个顺手的炉子。作为 AI 工程师日常必不可缺的「炼丹炉」,「PyTorch 还是 TensorFlow?」已成为知乎、Reddit 等炼丹师出没之地每年都会讨论的热门话题。
业界流传一种说法:PyTorch 适合学术界,TensorFlow 适合工业界。毕竟,PyTorch 是用户最喜欢的框架,API 非常友好,Eager 模式让模型搭建和调试过程变得更加容易,不过,它的静态图编译和部署体验还不令人满意。TensorFlow 恰恰相反,静态编译和部署功能很完备,不过其调试体验让人欲哭无泪。
那么问题来了:鱼和熊掌真的不可兼得吗?未必,来自北京的一流科技团队推出的开源深度学习框架 OneFlow 已经做到了。
等等,OneFlow 一直主打分布式和高性能,易用性也能和 PyTorch一样吗?听说过 OneFlow 的人一定会发出这样的疑问。
没错,从 2016 年底立项之日起,OneFlow 就是为大规模分布式而生,特色之一就是静态图机制,2020 年 7 月在 GitHub 上开源时还不支持动态图。不过,OneFlow 团队用一年多时间自研了动态图引擎, OneFlow v0.7.0 版本已支持和 PyTorch 一模一样的 Eager 体验,也就是说,OneFlow 实现了同时支持动态图和静态图。不仅如此,OneFlow 编程 API 完全和 PyTorch 兼容,常见深度学习模型只需修改一行 import oneflow as torch 就可以把 PyTorch 写的模型在 OneFlow 上跑起来。
不妨先到 OneFlow 视觉模型库 flowvision 看一看:https://github.com/Oneflow-Inc/vision,这个模型库已经支持计算机视觉领域图像分类、分割和检测等方向的经典 SOTA 模型 (见下表),这些模型都可以通过 import torch as oneflow 或 import oneflow as torch 实现自由切换。
OneFlow 和 PyTorch 兼容之后,用户可以像使用 PyTorch 一样来使用 OneFlow ,对模型效果比较满意之后,可以继续使用 OneFlow 扩展到大规模分布式或使用静态图部署模型。听上去是不是 too good to be true?
在下面的案例中,一家头部通信公司基于 PyTorch 的业务模型快速方便地迁移成 OneFlow 的模型,并进行大幅度的训练/推理性能优化、部署上线,短短几天时间就让业务得以按时上线部署,且各项性能指标均大幅超出预期!
他们究竟是如何做到的?先从项目背景说起。
1
为什么选择 OneFlow?
因业务发展需求,这家通信公司近期将上线一款基于深度学习的图像识别应用,该项目的业务需求有如下五个特点:
import torchvision.models as models_torch
import 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 的收敛曲线也一致,在小数据集上的精度完全一致。class ResNet101Graph(oneflow.nn.Graph):
def __init__(self, input_shape, input_dtype=oneflow.float32):
super().__init__()
# 添加 ResNet101 nn.Module
self.model = ResNet101Module(input_shape, input_dtype)
self.loss_fn = ResNet101_loss_fn
# 添加 对应的 Optimizer
of_sgd = torch.optim.SGD(self.model.parameters(), lr=1.0, momentum=0.0)
self.add_optimizer(of_sgd)
# 配置静态图的自动优化选项
_config_graph(self)
def build(self, input):
# 类似 nn.Module 的 forward 方法,这里是构图,包括了构建后向图,所以叫 build
out = self.model(input)
loss = self.loss_fn(out)
# build 里面支持构建后向图
loss.backward()
return loss
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% 的加速。一般推理时只需要前向计算,后向计算是不需要的,但在用户这个特殊的模型里,部署和推理也是需要后向计算,只是不需要模型更新,这就导致用户写代码时为了保留后向计算也误把参数更新的逻辑保留下来了。据此可以省略参数的梯度计算,这里大概带来了 75% 的加速;
进而发现原任务(前向、后向、前向)中的第二次前向在部署时是多余的,可以裁剪掉,这里大概带来了大约 33% 的加速。
增大 batch_size,默认 batch_size 为 1,此时 GPU 利用率为 30%,当增大到 16 时,最高可以达到 90%,这里大约得到了 155% 的加速;
由于数据预处理在 CPU,网络计算在 GPU,两种设备接力执行,这时使用 2 进程进行,给数据加载部分加一个互斥锁,可以比较简易的实现 CPU 和 GPU 两级流水线,这里带来了 80% 的加速。
对于 ResNet101,batch_size 设置为 16,在 nn.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)
def _config_graph(graph):
if args.tensorrt:
# 使用 TensorRT 后端执行
graph.config.enable_tensorrt(True)
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")作为模型版本。模型就会通过 OneFlow 的 XRT 模块自动使用 TensorRT 进行推理,此外 OneFlow Serving 还支持类似的 “--enable-openvino”。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
curl -v localhost:8000/v2/health/ready
返回值为 HTTP/1.1 200 OK,表示服务正在正常工作。
#/usr/bin/env python3
import numpy as np
import tritonclient.http as httpclient
from 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 / 255
image = np.expand_dims(image, axis=0)
# Transpose NHWC to NCHW
image = 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(OneFlowTechnology)。
如有侵权,请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”,欢迎正在阅读的你也加入,一起分享。
|