torchpipe

Pytorch 内的多线程流水线并行库

English | 简体中文

torchpipe是介于底层加速库（如tensorrt，opencv，torchscript）以及 RPC（如thrift, gRPC）之间并与他们严格解耦的多实例流水线并行库；对外提供面向pytorch前端的线程安全函数接口，对内提供面向用户的细粒度后端扩展。

Torchpipe is a multi-instance pipeline parallel library that acts as a bridge between lower-level acceleration libraries (such as TensorRT, OpenCV, TorchScript) and RPC frameworks (like Thrift, gRPC), ensuring a strict decoupling from them. It offers a thread-safe function interface for the PyTorch frontend at a higher level, while empowering users with fine-grained backend extension capabilities at a lower level.

注意事项

版本说明：推荐使用最新tag以及对应release
master分支用于发布版本更新，develop分支用于提交代码和日常开发；

快速开始

1. 安装

参见安装文档

2. 获取恰当的模型文件(目前支持 onnx, trt engine等)

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") 
resnet18 = models.resnet18(pretrained=True).eval().cuda()
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}")

3. 现在你可以并发调用单模型了

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)  # 失败则此键值一定不存在，即使输入时已经存在。

4. 我们的核心功能为多个节点间的一系列流水线设施。

更多信息，访问 Torchpipe的文档。

5. RoadMap

torchpie目前处于一个快速迭代阶段，我们非常需要你的帮助。相对于贡献形式，我们对内容方面更为在意。欢迎通过issues或者merge requests等方式进行反馈。贡献指南。

我们的最终目标是让服务端高吞吐部署尽可能简单。为了实现这一目标，我们将积极自我迭代，也愿意参与有相近目标的其他项目。

2023年度和2024年度 RoadMap

技术报告
大模型方面的示例
公开的基础镜像和pypi(manylinux)
优化编译系统，分为core,pplcv,model/tensorrt,opencv等模块
基础结构优化。包含python与c++交互，异常，日志系统，跨进程后端的优化；

潜在未完成的研究方向

单节点调度和多节点调度后端，他们与计算后端无本质差异，需要更多面向用户进行解耦，我们想要将这部分优化为用户API的一部分；
针对多节点的调试工具。由于在多节点调度中，使用了模拟栈设计，比较容易设计节点级别的调试工具；
负载均衡

6. 致谢

我们的代码库使用或者修改后使用了多个开源库，请查看致谢了解更多详细信息。

Alwaysssssss/torchpipe