An Alternative for Triton Inference Server. Boosting DL Service Throughput 1.5-4x by Ensemble Pipeline Serving with Concurrent CUDA Streams for PyTorch/LibTorch Frontend and TensorRT/CVCUDA, etc., Backends
English | 简体中文
torchpipe是 介于底层加速库(如tensorrt,opencv,CVCUDA, ppl.cv)以及 RPC(如thrift, gRPC)之间并与他们严格解耦的多实例流水线并行库;对外提供面向pytorch前端的线程安全函数接口,对内提供面向用户的细粒度后端扩展。
torchpipe是 Triton Inference Server 的一个代替选择,主要功能类似于其共享显存,Ensemble, BLS机制。
生产级别:在网易智企内部,每天有海量调用由Torchpipe支持。
tensorrt在max_batch_size=4时,很多时候输入1张和4张时结果有差异,这是正常的。但是此时固定输入只有有限种类(一般为2)结果
git clone https://github.com/torchpipe/torchpipe.git
cd torchpipe/ && git submodule update --init --recursive
export img_name=nvcr.io/nvidia/pytorch:22.12-py3
# or build docker image by yourself (recommend, for tensorrt 9.3):
# docker build --network=host -f ./docker/Dockerfile -t trt-9 thirdparty/
# export img_name=trt-9
docker run --rm --gpus=all --ipc=host --network=host -v `pwd`:/workspace --shm-size 1G --ulimit memlock=-1 --ulimit stack=67108864 --privileged=true -w/workspace -it $img_name /bin/bash
python setup.py install
cd examples/resnet18 && python resnet18.py
详见 安装文档
import torchvision.models as models
resnet18 = models.resnet18(pretrained=True).eval().cuda()
import tempfile, os, torch
model_path = os.path.join(tempfile.gettempdir(), "./resnet18.onnx")
data_bchw = torch.rand((1, 3, 224, 224)).cuda()
print("export: ", model_path)
torch.onnx.export(resnet18, data_bchw, model_path,
opset_version=17,
do_constant_folding=True,
input_names=["in"], output_names=["out"],dynamic_axes={"in": {0: "x"},"out": {0: "x"}})
# os.system(f"onnxsim {model_path} {model_path}")
import torch, torchpipe
model = torchpipe.pipe({'model': model_path,
'backend': "Sequential[cvtColorTensor,TensorrtTensor,SyncTensor]", # 后端引擎, 可见后端API参考文档。
'instance_num': 2, 'batching_timeout': '5', # 实例数和超时时间
'max': 4, # 模型优化范围最大值,也可以为 '4x3x224x224'
'mean': '123.675, 116.28, 103.53',#255*"0.485, 0.456, 0.406",
'std': '58.395, 57.120, 57.375', # 将融合进tensorrt网络中
'color': 'rgb'}) # cvtColorTensor后端的参数: 目标颜色空间顺序
data = torch.zeros((1, 3, 224, 224)) # or torch.from_numpy(...)
input = {"data": data, 'color': 'bgr'}
model(input) # 可多线程并行调用
# 使用 "result" 作为数据输出标识;当然,其他键值也可自定义写入
print(input["result"].shape) # 失败则此键值一定不存在,即使输入时已经存在。
纯c++ API 可通过 [libtorch+cmake] 或者 [pybind11]的方式获得.
更多信息,访问 Torchpipe的文档 。
torchpie目前处于一个快速迭代阶段,我们非常需要你的帮助。欢迎通过issues或者merge requests等方式进行反馈。贡献指南。
我们的最终目标是让服务端高吞吐部署尽可能简单。为了实现这一目标,我们将积极自我迭代,也愿意参与有相近目标的其他项目。
近期 RoadMap
潜在未完成的研究方向
我们的代码库使用或者修改后使用了多个开源库,请查看致谢了解更多详细信息。