兼容PyTorch,25倍加速,国产!
共 11856字,需浏览 24分钟
·
2022-05-12 23:44
↑ 关注 + 星标 ,每天学Python新技能
后台回复【大礼包】送你Python自学大礼包
数据量大:数据库中有过亿级别的图片
模型简单:比较常规的分类模型
400 多张显卡,短期内无法扩容
对于训练/推理的吞吐有硬性指标
上线时间紧迫
1、OneFlow 是众多深度学习框架中,API 与 PyTorch 兼容性最高的,这样方便工程师用最少的时间/人力成本,对已有项目代码进行迁移,减少学习成本。
2、OneFlow 动静转换十分方便,动态图(Eager)模式的代码简单改动几行就能转换为静态图(nn.Graph)模式。
3、OneFlow 在框架层面做了大量优化,nn.Graph 提供了简洁、丰富的性能优化选项,如算子融合(Kernel Fusion)、自动混合精度训练 (Auto Mixed Precision Training) 等。
将已有 PyTorch 的项目代码完全迁移到 OneFlow
将项目代码由动态图模式(Eager Mode)改造为静态图模式(Graph Mode)
开启 OneFlow Graph 模式下的各种优化选项并训练模型
用 Serving 模块部署模型上线
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)
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
resnet101_graph = ResNet101Graph((args.batch_size, 3, img_shape[1], img_shape[0]))
for i in range(m):
loss = resnet101_graph(images)
一般推理时只需要前向计算,后向计算是不需要的,但在用户这个特殊的模型里,部署和推理也是需要后向计算,只是不需要模型更新,这就导致用户写代码时为了保留后向计算也误把参数更新的逻辑保留下来了。据此可以省略参数的梯度计算,这里大概带来了 75% 的加速;
进而发现原任务(前向、后向、前向)中的第二次前向在部署时是多余的,可以裁剪掉,这里大概带来了大约 33% 的加速。
增大 batch_size,默认 batch_size 为 1,此时 GPU 利用率为 30%,当增大到 16 时,最高可以达到 90%,这里大约得到了 155% 的加速;
由于数据预处理在 CPU,网络计算在 GPU,两种设备接力执行,这时使用 2 进程进行,给数据加载部分加一个互斥锁,可以比较简易的实现 CPU 和 GPU 两级流水线,这里带来了 80% 的加速。
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)
打开混合精度,测试得到了 36% 的加速
再打开卷积试跑优化,测试得到了 7% 的加速,总加速为 43%
再打开 pad 和 conv 算子融合,测试得到了 19% 的加速,总加速为 62%
再打开 add 的算子的融合,测试得到了 2% 的加速,总加速为 64%
def _config_graph(graph):
if args.tensorrt:
# 使用 TensorRT 后端执行
graph.config.enable_tensorrt(True)
nn.Graph 的使用教程,https://docs.oneflow.org/en/master/basics/08_nn_graph.html
nn.Graph 的 API 文档,https://oneflow.readthedocs.io/en/master/graph.html
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()
unused_output = inference_graph(flow.zeros(1, 3, 224, 224))
flow.save(inference_graph, "model")
docker run --rm --runtime=nvidia --network=host -v$(pwd)/model:/models/resnet101/1 \
oneflowinc/oneflow-serving:nightly
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
#/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项目地址:https://github.com/Oneflow-Inc/oneflow/
OneFlow用户文档:https://docs.oneflow.org/master/index.html
推荐阅读