构建Torch-TensorRT的C++ 库
简介
Torch-TensorRT是一个将PyTorch模型转化为TensorRT可执行引擎的库,旨在将PyTorch模型的推理性能提升至接近TensorRT的水平。它充分利用了TensorRT在优化CUDA内核和执行计划方面的优势,使得PyTorch模型能够在NVIDIA GPU上实现高效的推理。Torch-TensorRT简化了PyTorch到TensorRT的转换过程,降低了使用门槛,使得开发者能够轻松地将PyTorch模型部署到生产环境中,实现高性能的实时推理。这一库的出现,极大地推动了PyTorch在深度学习推理领域的应用和发展。
安装环境
#先建一个虚拟环境,并进入
conda create -n Torch-TensorRT python=3.8 -y
conda activate Torch-TensorRT
安装 Torch-TensorRT
#安装自己对应环境cuda版本的Torch-TensorRT,这里是cuda11.6
pip install torch torch-tensorrt tensorrt --extra-index-url https://download.pytorch.org/whl/cu116
安装 Bazel
参考这里:Ubuntu安装Bazel
下载源码
git clone https://github.com/pytorch/TensorRT.git
cd TensorRT
修改WORKSPACE 文件
- 这里修改为刚才下载的torch-tensorrt包的位置
- 修改cuda路径
- 修改libtorch路径
- 修改libtorch_pre_cxx11_abi的路径
- 在英伟达官网下载对应自己环境的cudnn并计算校验和
使用 sha256sum 命令来计算文件的校验和
- 修改cudnn部分的内容
- 同样的方法修改tensorrt部分的内容
构建 C++ 库
- 若要构建release的库
bazel build //:libtorchtrt -c opt
- 若要构建debug的库
bazel build //:libtorchtrt -c dbg
- 若要构建pre_cxx11_abi的库
bazel build //:libtorchtrt --config pre_cxx11_abi -c [dbg/opt]
成功后就会生成找到包含文件和库的压缩包bazel-bin