0x16 - TensorRT 使用之环境搭建

发布： 更新于 分类： 阅读时长 ≈ 2 分钟

TensorRT

Nvidia 的 TensorRT 是一套专注于为深度学习加速的推理开发框架，面向的自然是机器学习相关的应用。
而能够运用这套推理框架的，自然也是 Nvidia 的 GPU 了。

将 ONNX 或 libtorch 的模型导出为 TRT-Engine 以后，这套开发框架能为深度学习应用提供可观的速度优化。
就此告别推理速度过慢的烦恼。

本文旨在建立一套可以根据需求配置的 TensorRT 开发环境，
具体在部署阶段如何精简库文件需要根据自己的需求来决定。

官方 Github：TensorRT Open Source Software

Docker 环境搭建

这里默认大家都拥有了 Nvidia 相关的推理硬件，相关的操作系统。
我们以搭建可用的 docker 环境为目标。

官方的 docker 可以去 docker hub 或者 官方 github 寻找。
考虑到官方的 docker 存在一些兼容性的问题，这里提供一种自己搭建 docker 镜像的方法。

0x00 事前准备

  • 准备好 Nvidia Cuda 环境和 Nvidia Docker，或是 Nvidia Container Toolkits。

0x01 Dockerfile

官方的 dockerfile 是可用的，如果官方的环境需求能满足的话，直接使用 docker build 即可。
以 TRT 8.6 的官方 dockerfile 为例 https://github.com/NVIDIA/TensorRT/tree/release/8.6/docker
这里对如何修改 CUDA 版本以及 TRT 版本作一些简单的说明。

以 ubuntu 20.04 这个版本为准，以下为官方 dockerfile，请查看其中的注释。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
# 这里可以修改具体的 CUDA 版本
# 考虑到使用的是官方的 Cuda 镜像，可以去这个网址看看具体版本号应该怎么填写
# https://hub.docker.com/r/nvidia/cuda/tags
ARG CUDA_VERSION=12.1.1

# 官方镜像没有配置镜像源，直接使用可能安装速度会十分缓慢
# 这里推荐先将镜像的apt源修改过后，使用编辑过的镜像来加快镜像构建速度
FROM nvidia/cuda:${CUDA_VERSION}-cudnn8-devel-ubuntu20.04
LABEL maintainer="NVIDIA CORPORATION"

ENV TRT_VERSION 8.6.1.6
SHELL ["/bin/bash", "-c"]

# Setup user account
ARG uid=1000
ARG gid=1000
RUN groupadd -r -f -g ${gid} trtuser && useradd -o -r -l -u ${uid} -g ${gid} -ms /bin/bash trtuser
RUN usermod -aG sudo trtuser
RUN echo 'trtuser:nvidia' | chpasswd
RUN mkdir -p /workspace && chown trtuser /workspace

# Required to build Ubuntu 20.04 without user prompts with DLFW container
ENV DEBIAN_FRONTEND=noninteractive

# Update CUDA signing key
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub

# Install requried libraries
RUN apt-get update && apt-get install -y software-properties-common
RUN add-apt-repository ppa:ubuntu-toolchain-r/test
RUN apt-get update && apt-get install -y --no-install-recommends \
    libcurl4-openssl-dev \
    wget \
    git \
    pkg-config \
    sudo \
    ssh \
    libssl-dev \
    pbzip2 \
    pv \
    bzip2 \
    unzip \
    devscripts \
    lintian \
    fakeroot \
    dh-make \
    build-essential

# Install python3
RUN apt-get install -y --no-install-recommends \
      python3 \
      python3-pip \
      python3-dev \
      python3-wheel &&\
    cd /usr/local/bin &&\
    ln -s /usr/bin/python3 python &&\
    ln -s /usr/bin/pip3 pip;

# Install TensorRT
RUN if [ "${CUDA_VERSION}" = "10.2" ] ; then \
    v="${TRT_VERSION%.*}-1+cuda${CUDA_VERSION}" &&\
    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub &&\
    apt-get update &&\
    sudo apt-get install libnvinfer8=${v} libnvonnxparsers8=${v} libnvparsers8=${v} libnvinfer-plugin8=${v} \
        libnvinfer-dev=${v} libnvonnxparsers-dev=${v} libnvparsers-dev=${v} libnvinfer-plugin-dev=${v} \
        python3-libnvinfer=${v} libnvinfer-dispatch8=${v} libnvinfer-dispatch-dev=${v} libnvinfer-lean8=${v} \
        libnvinfer-lean-dev=${v} libnvinfer-vc-plugin8=${v} libnvinfer-vc-plugin-dev=${v} \
        libnvinfer-headers-dev=${v} libnvinfer-headers-plugin-dev=${v}; \
else \
    ver="${CUDA_VERSION%.*}" &&\
    if [ "${ver%.*}" = "12" ] ; then \
        ver="12.0"; \
    fi &&\
    v="${TRT_VERSION}-1+cuda${ver}" &&\
    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub &&\
    apt-get update &&\
    sudo apt-get -y install libnvinfer8=${v} libnvonnxparsers8=${v} libnvparsers8=${v} libnvinfer-plugin8=${v} \
        libnvinfer-dev=${v} libnvonnxparsers-dev=${v} libnvparsers-dev=${v} libnvinfer-plugin-dev=${v} \
        python3-libnvinfer=${v} libnvinfer-dispatch8=${v} libnvinfer-dispatch-dev=${v} libnvinfer-lean8=${v} \
        libnvinfer-lean-dev=${v} libnvinfer-vc-plugin8=${v} libnvinfer-vc-plugin-dev=${v} \
        libnvinfer-headers-dev=${v} libnvinfer-headers-plugin-dev=${v}; \
fi

# Install PyPI packages
# 这里安装 Python3 的依赖可能会比较慢，同样的，可以加入第三方源来加速构建

RUN pip3 install --upgrade pip
# RUN pip3 install --upgrade pip -i https://mirrors.bfsu.edu.cn/pypi/web/simple

RUN pip3 install setuptools>=41.0.0
# RUN pip3 install setuptools>=41.0.0 -i https://mirrors.bfsu.edu.cn/pypi/web/simple

COPY requirements.txt /tmp/requirements.txt

RUN pip3 install -r /tmp/requirements.txt
# 这一步由于需要借助 torch 的官方源安装，十分缓慢，可以考虑手动提前安装部分包
# 这里提前安装了 torch 以及 torchvision
# 具体版本号可以根据安装时的超时报错决定
# 以下两个 whl 包需要先手动下载
# 网址可以参考 https://download.pytorch.org/whl/cu113/torch_stable.html
# COPY torch-1.10.2+cu113-cp38-cp38-linux_x86_64.whl /tmp/torch-1.10.2+cu113-cp38-cp38-linux_x86_64.whl
# RUN pip3 install /tmp/torch-1.10.2+cu113-cp38-cp38-linux_x86_64.whl -i https://mirrors.bfsu.edu.cn/pypi/web/simple
# COPY torchvision-0.11.3+cu113-cp38-cp38-linux_x86_64.whl /tmp/torchvision-0.11.3+cu113-cp38-cp38-linux_x86_64.whl
# RUN pip3 install /tmp/torchvision-0.11.3+cu113-cp38-cp38-linux_x86_64.whl -i https://mirrors.bfsu.edu.cn/pypi/web/simple
# RUN pip3 install -r /tmp/requirements.txt -i https://mirrors.bfsu.edu.cn/pypi/web/simple

RUN pip3 install jupyter jupyterlab
# RUN pip3 install jupyter jupyterlab -i https://mirrors.bfsu.edu.cn/pypi/web/simple
# Workaround to remove numpy installed with tensorflow

RUN pip3 install --upgrade numpy
# RUN pip3 install --upgrade numpy -i https://mirrors.bfsu.edu.cn/pypi/web/simple

# Install Cmake
RUN cd /tmp && \
    wget https://github.com/Kitware/CMake/releases/download/v3.14.4/cmake-3.14.4-Linux-x86_64.sh && \
    chmod +x cmake-3.14.4-Linux-x86_64.sh && \
    ./cmake-3.14.4-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
    rm ./cmake-3.14.4-Linux-x86_64.sh

# Download NGC client
RUN cd /usr/local/bin && wget https://ngc.nvidia.com/downloads/ngccli_cat_linux.zip && \
    unzip ngccli_cat_linux.zip && chmod u+x ngc-cli/ngc && \
    rm ngccli_cat_linux.zip ngc-cli.md5 && echo "no-apikey\nascii\n" | ngc-cli/ngc config set

# Set environment and working directory
ENV TRT_LIBPATH /usr/lib/x86_64-linux-gnu
ENV TRT_OSSPATH /workspace/TensorRT
ENV PATH="${PATH}:/usr/local/bin/ngc-cli"
ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${TRT_OSSPATH}/build/out:${TRT_LIBPATH}"
WORKDIR /workspace

# 这里设置了 trtuser 用户以避免一些错误更改导致环境出现问题
# 使用 root 也没有问题
USER trtuser
RUN ["/bin/bash"]

0x02 搭建步骤

1
2
3
4
5
6
7
8
9
10
11
git clone https://github.com/NVIDIA/TensorRT.git

cd tensorrt
git submodule update --init --recursive

vi docker/my.Dockerfile
# write your own dockerfile
# an example of cu118_devel
docker build -f docker/my.Dockerfile --tag tensorrt:cu118_devel --network host .

# 期间可能会有网络问题导致安装中断，换源或使用离线安装包都可以是解决方法

0x03 启动环境

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
#!/bin/bash
nvidia-docker run -it \
--gpus all \
--shm-size 256G \
--name temp_trt_cu118 \
--network host \
-e NVIDIA_DRIVER_CAPABILITIES=video,compute,utility \
tensorrt:cu118_devel \
/bin/bash

# 如果想使用默认的 root 权限，可以增加 --user root 这个选项启用

# 注意要添加外部路径到 docker 内部的 workspace 以方便后续编译使用
# 考虑到 TRT 强硬件关联性，这里以 docker cp 而非 volume 挂载的方式以避免奇怪的问题
docker cp /path/to/TensorRT temp_trt_cu118:/workspace/TensorRT

注意要添加外部路径到 docker 内部的 workspace 以方便后续使用。
至此 TensorRT 的环境算是搞定了。

Bazel 支持

要使 bazel 支持 TensorRT 的依赖，可以看以下方法

0x00 安装 Bazel

0x01 编译 TensorRT Open Source Software

使用官方的 Cmake 方式可以编译一些版本不一致的，可供测试和模型导出的 binary 文件。
具体参数配置选项可以参考官方文档 https://github.com/NVIDIA/TensorRT

这里给出一个编译的方式

1
2
3
4
5
6
7
8
9
10
11
# 如果你没有准备好 TensorRTOSS
# 考虑到 TensorRT 比较依赖具体的硬件环境，推荐每次在新设备上部署开发环境都做一下以下步骤

# 如果已经在外部 init 过
docker cp /path/to/TensorRT temp_trt_cu118:/workspace/TensorRT

# 也可以从零开始
cd /workspace
git clone https://github.com/NVIDIA/TensorRT.git
cd TensorRT
git submodule update --init --recursive

如果要用 trtexec 指令或添加自定义算子，或许需要编译以下内容

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
cd /workspace/TensorRT
mkdir -p build && cd build

cmake .. \
  -DTRT_LIB_DIR=$TRT_LIBPATH \
  -DTRT_OUT_DIR=`pwd`/out \
  -DCMAKE_BUILD_TYPE=release \
  -DCUDA_VERSION=11.8.0 \
  -DCUDNN_VERSION=8.8 \
  -DGPU_ARCHS="70 75 80"

# 需要注意的是，这里的 cuda 版本和 cudnn 版本还有 GPU archs 需要和你实际情况保持一致
# 否则会出现运行不起来的情况

make -j$(nproc)

结束以后可以在 out 文件夹下找到静态库和可执行文件。

0x02 TensorRT BUILD 文件编写

需要首先在 TensorRT OSS 文件夹下建立至 cuda、drivers 的软链并新建 BUILD 文件

1
2
3
cd /workspace/tensorRT
ln -s /usr/local/cuda-11.8/ ./cuda
ln -s /usr/lib/x86_64-linux-gnu/ ./trt_lib

如果使用的是 TRT 9.1 或以上的话，可能会需要以下的包，额外安装一下。

1
2
3
4
extra: TRT 9.1+ apt 包

apt-get install libmlx5-1
apt-get install -y rdmacm-utils

文件夹下新建一个 BUILD 文件，内容如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
# BUILD FILE
package(
    default_visibility = [
        "//visibility:public",
    ],
)

cc_library(
    name = "TensorRT",
    srcs = [
        "trt_lib/libnvinfer.so",
        "trt_lib/libnvcaffe_parser.so",
        "trt_lib/libnvinfer_plugin.so",
        "trt_lib/libnvonnxparser.so",
        "trt_lib/libnvparsers.so"
    ],
    hdrs = glob([
        "include/*.h",
    ]),
    includes = [
        "include",
        "cuda/include"
    ],
    deps = [
        ":cuda_envs"
    ]
)

cc_library(
    name = "cuda_envs",
    srcs = glob([
        "cuda/lib64/*.so",
        "cuda/lib64/stub/*.so"
    ]),
    hdrs = glob([
        "cuda/include/*.h",
        "cuda/include/**/*.h",
        "cuda/include/**/*.hpp"
    ])
)

至此，若想依赖 TensorRT 的环境和头文件，引用这个 TensorRT 即可。