项目概述
tiny-tensorrt 是一个简单易用的 NVIDIA TensorRT 封装工具,专为 ONNX 模型设计,并提供 C++ 和 Python 两种 API 接口。用户可以通过极少的代码行数,快速地将模型部署到生产环境中。
项目背景
tiny-tensorrt 是在更好的引擎构建替代方案出现之前的一个有益尝试。不过,随着 TensorRT 的 Python API 和 trtexec/polygraphy 工具的推广,tiny-tensorrt 项目已不再维护。因此,使用者可以选择这些新工具快速构建引擎。
功能特色
使用 tiny-tensorrt,用户可以轻松进行以下操作:
- 创建 TensorRT 引擎。
- 在主机与设备之间传输数据。
- 高效地执行前向传播运算。
下面是一个典型的 C++ 使用示例代码:
Trt* net = new Trt();
net->SetFP16();
net->BuildEngine(onnxModel, engineFile);
net->CopyFromHostToDevice(input, inputBindIndex);
net->Forward();
net->CopyFromDeviceToHost(output, outputBindIndex);
安装指南
为了使用 tiny-tensorrt,用户需预先安装以下库:
- CUDA
- CUDNN
- TensorRT
支持的 CUDA 版本包括:10.2, 11.0, 11.1, 11.2, 11.3 和 11.4。支持的 TensorRT 版本包括:7.0, 7.1, 7.2, 8.0, 8.2 和 8.4。
安装步骤如下:
-
更新软件包并安装构建工具:
sudo apt-get update -y sudo apt-get install cmake zlib1g-dev
-
如果需要 Python 绑定,安装 Python 依赖:
sudo apt-get install python3 python3-pip pip3 install numpy
-
克隆项目仓库:
git clone --recurse-submodules -j8 https://github.com/zerollzeng/tiny-tensorrt.git
-
构建项目:
cd tiny-tensorrt mkdir build && cd build cmake .. && make
用户构建完成后,可以将库文件 libtinytrt.so
和头文件 Trt.h
集成到自己的项目中,并可使用 Python 模块 pytrt.so
。
文档与许可证
关于 tiny-tensorrt 的更多信息,请参考项目的 Wiki 页面。同时,需要注意的是,任何第三方模块和 TensorRT 的使用需遵循其各自的许可证协议。至于项目中由作者编写的部分,用户可以随意使用。