构建Torch-TensorRT的C++ 库

构建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
python -m 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 文件

1. 这里修改为刚才下载的torch-tensorrt包的位置
   
2. 修改cuda路径
   
3. 修改libtorch路径
   
4. 修改libtorch_pre_cxx11_abi的路径
   
5. 在英伟达官网下载对应自己环境的cudnn并计算校验和
   使用 sha256sum 命令来计算文件的校验和
   
6. 修改cudnn部分的内容
   
7. 同样的方法修改tensorrt部分的内容
   

构建 C++ 库

1. 若要构建release的库

bazel build //:libtorchtrt -c opt

2. 若要构建debug的库

bazel build //:libtorchtrt -c dbg

3. 若要构建pre_cxx11_abi的库

bazel build //:libtorchtrt --config pre_cxx11_abi -c [dbg/opt]

成功后就会生成找到包含文件和库的压缩包bazel-bin