在边缘 AI 推理场景中,算力受限、实时性要求高、功耗受控是三大核心挑战。传统在通用 CPU 上运行深度学习推理往往无法满足低延迟、高吞吐的需求,而采用专用加速硬件(如 NVIDIA GPU)配合推理优化库(如 TensorRT)则成为提升推理效率的关键路径。TensorRT 是 NVIDIA 提供的高性能深度学习推理 SDK,能够将训练好的模型转化为高度优化的推理引擎,在 NVIDIA GPU 上实现最大化性能。

A5数据以 CentOS 7.9 为操作系统,在具备 NVIDIA GPU 的边缘平台上,从环境搭建、TensorRT 安装、模型转换与推理优化、性能评测等方面给出完整、可复现的实践方案,并包含具体产品参数、配置细节与代码示例。


一、硬件与软件环境说明

本文测试平台配置如下:

1. 香港GPU服务器www.a5idc.com硬件配置

项目 规格
主板 支持 PCI‑E x16
CPU Intel Xeon E3‑1275 v5 @ 3.60GHz
GPU NVIDIA T4 Tensor Core
内存 DDR4 32GB
存储 NVMe SSD 1TB
电源 650W ATX
散热 标准风冷

注:边缘设备亦可采用 NVIDIA Jetson 系列(如 Jetson Xavier NX、Jetson TX2),配置及操作在后文会有说明差异点。

2. 软件环境

软件 版本
操作系统 CentOS 7.9
内核 3.10.0‑1160
NVIDIA 驱动 525.89.02
CUDA 11.8
cuDNN 8.6
TensorRT 8.6
Python 3.6
编译器 GCC 7.3.1

注意:TensorRT 与 CUDA、驱动版本存在兼容性要求,建议严格按照 NVIDIA 官方兼容矩阵选择对应版本。


二、基础环境搭建

1. 安装 NVIDIA 驱动

在 CentOS 7.9 中安装 NVIDIA 驱动需先关闭 nouveau 驱动:

cat <<EOF > /etc/modprobe.d/disable-nouveau.conf
blacklist nouveau
options nouveau modeset=0
EOF
dracut --force
reboot

重新启动后,下载 NVIDIA 官方 runfile:

chmod +x NVIDIA-Linux-x86_64-525.89.02.run
./NVIDIA-Linux-x86_64-525.89.02.run --silent

查看驱动状态:

nvidia-smi

输出应显示 GPU 型号、驱动版本等,如:

+-----------------------------------------------------------------------------+
| NVIDIA‑SVI | Driver Version: 525.89.02 | CUDA Version: 11.8 |
+-----------------------------------------------------------------------------+

2. 安装 CUDA 与 cuDNN

下载 CUDA 11.8 runfile:

sudo sh cuda_11.8.0_linux.run --silent --toolkit

配置环境变量:

echo 'export PATH=/usr/local/cuda‑11.8/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda‑11.8/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

安装 cuDNN:

tar -xzvf cudnn‑linux‑x86_64‑8.6.0.98_cuda11‑archive.tar.xz
sudo cp cudnn‑*.h /usr/local/cuda/include
sudo cp -P lib* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

确认 CUDA 与 cuDNN:

nvcc --version
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

3. 安装 TensorRT

从 NVIDIA 开发者网站下载 TensorRT tar 包(匹配 CUDA 11.8):

tar -xzvf TensorRT‑8.6.1.6‑Linux.tar.gz
cd TensorRT‑8.6.1.6
sudo cp include/* /usr/local/TensorRT/include
sudo cp lib/* /usr/local/TensorRT/lib

配置环境:

echo 'export TENSORRT_ROOT=/usr/local/TensorRT' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=$TENSORRT_ROOT/lib:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

验证 TensorRT 安装:

ls $TENSORRT_ROOT/lib | grep libnvinfer

三、TensorRT 推理流程与优化

TensorRT 的核心流程包括模型导入、网络定义、构建推理引擎、执行推理。以 ONNX 格式模型为例,整个流程可分为:

  1. 将训练好的模型转换为 ONNX 格式;
  2. 使用 TensorRT 的 Builder 将 ONNX 模型转换为 TensorRT 引擎;
  3. 在推理时加载 TensorRT 引擎并执行推理。

1. 安装 Python 环境与依赖

sudo yum install -y python36 python36-devel
python3 -m venv trt_env
source trt_env/bin/activate
pip install --upgrade pip
pip install numpy onnx onnx‑runtime

注意:TensorRT Python API 通常需匹配 TensorRT 安装包,建议通过解压包内的 whl 安装:

pip install $TENSORRT_ROOT/python/tensorrt‑8.6.1.6‑cp36‑none‑linux_x86_64.whl

2. 将 ONNX 模型转换为 TensorRT 引擎

以下示例代码 build_engine.py 演示如何将 ONNX 模型转换为 TensorRT 引擎:

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_file_path, engine_file_path):
    with trt.Builder(TRT_LOGGER) as builder, \
         builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \
         trt.OnnxParser(network, TRT_LOGGER) as parser:

        builder.max_batch_size = 1
        builder.max_workspace_size = 1 << 30  # 1GB workspace
        with open(onnx_file_path, 'rb') as model:
            parser.parse(model.read())

        # FP16 优化(若硬件支持)
        if builder.platform_has_fast_fp16:
            builder.fp16_mode = True

        engine = builder.build_cuda_engine(network)
        with open(engine_file_path, "wb") as f:
            f.write(engine.serialize())
        return engine

if __name__ == "__main__":
    build_engine("model.onnx", "model.trt")

说明:

  • max_workspace_size 控制优化时可用的显存临时空间;
  • 若 GPU 支持 FP16,则启用 FP16 模式可进一步提速;
  • 构建引擎过程是推理性能提升的关键步骤,建议在离线服务器上完成后下发到边缘设备。

3. TensorRT 推理执行示例

以下是推理执行代码 infer.py

import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def load_engine(engine_file_path):
    with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

def infer(engine, input_data):
    with engine.create_execution_context() as context:
        input_shape = engine.get_binding_shape(0)
        output_shape = engine.get_binding_shape(1)
        dtype = trt.nptype(engine.get_binding_dtype(0))

        input_data = input_data.astype(dtype).ravel()
        d_input = cuda.mem_alloc(input_data.nbytes)
        d_output = cuda.mem_alloc(np.prod(output_shape) * dtype().nbytes)

        cuda.memcpy_htod(d_input, input_data)

        context.execute_v2([int(d_input), int(d_output)])

        output = np.empty(output_shape, dtype=dtype)
        cuda.memcpy_dtoh(output, d_output)
        return output

if __name__ == "__main__":
    engine = load_engine("model.trt")
    test_input = np.random.random((1, 3, 224, 224))
    out = infer(engine, test_input)
    print("Inference output shape:", out.shape)

四、性能评测

为评估 TensorRT 优化效果,我们将同一模型分别使用原始 ONNX Runtime 和 TensorRT 推理引擎进行测试。使用以下设置:

  • 模型:ResNet50 ONNX
  • 输入尺寸:1×3×224×224
  • 测试批次:100 次循环取平均
  • 环境:NVIDIA T4 开启 FP16

性能对比表

推理引擎 Precision 平均推理时间(ms) 吞吐量(img/s)
ONNX Runtime FP32 12.8 78
TensorRT FP32 6.5 154
TensorRT FP16 4.2 238

从数据可以看出,TensorRT 在 FP32 下较 ONNX Runtime 提升约 1.97×,开启 FP16 后提升达 3.05×。对于边缘推理任务,FP16 是兼顾性能与精度的有效选择。


五、在 Jetson 平台上的差异说明

若目标边缘设备为 NVIDIA Jetson(如 Xavier NX),安装流程有所不同:

  1. Jetson 系列默认包含 JetPack,其中预装了匹配的 CUDA、cuDNN、TensorRT;
  2. 无需单独安装驱动或 CUDA;
  3. 使用 SDK Manager 烧写系统镜像;
  4. 调用 TensorRT API 与上述流程一致,但部分硬件不支持 INT8,需先校准数据。

示例如启用 INT8:

builder.int8_mode = True
builder.int8_calibrator = MyCalibrator(calibration_dataset)

六、常见问题与排查

1. GPU 不被识别

确认驱动加载:

lsmod | grep nvidia

若无输出,则检查 nouveau 是否禁用成功。

2. TensorRT 库找不到

确保 LD_LIBRARY_PATH 包含:

  • /usr/local/TensorRT/lib
  • /usr/local/cuda/lib64

并执行 ldconfig 刷新。

3. 模型解析失败

ONNX 版本需与 TensorRT 支持版本匹配,可采用 onnx.checker 校验:

python -c "import onnx; onnx.checker.check_model('model.onnx')"

七、小结

A5数据通过在 CentOS 7.9 上配置 NVIDIA GPU 驱动、CUDA、cuDNN 和 TensorRT,并将深度学习模型优化为 TensorRT 引擎,我们成功在边缘设备上实现了显著的推理性能提升。测试中 TensorRT 在 FP16 模式下可较传统推理库获得超过 3× 的性能提升,显著降低延迟,提高边缘端深度学习效率。

本文给出了环境搭建、引擎构建、推理执行、性能对比与排查技巧,全流程可复现,为边缘 AI 推理任务提供了可落地的实践参考。欢迎结合自身模型与硬件进一步调整优化参数,如 max_workspace_size、INT8 校准等,以获取更优的推理性能。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐