第2章：开发环境搭建：Ubuntu + Docker + ONNX Runtime + TensorRT 环境配置与验证

说实话，环境搭建这事儿，看着简单，但坑是真不少。

我见过太多同学，模型写好了，代码调通了，结果卡在环境配置上，一卡就是两三天。所以这一章，咱们把环境一次性搞定。我会把我在项目中踩过的坑、绕过的路，都给你讲清楚。

2.1 为什么选这套组合？

你可能会问：Ubuntu、Docker、ONNX Runtime、TensorRT，这四个东西凑一块儿，到底图啥？

我简单解释一下：

  • Ubuntu：嵌入式开发的主流系统。说白了，大部分AI芯片的SDK，都是基于Ubuntu的。你用Windows？嗯，也不是不行，但你会发现自己成了“孤岛”。
  • Docker：环境隔离的神器。我早期做项目时，经常因为系统库版本冲突，搞得焦头烂额。Docker一上，这些问题全没了。
  • ONNX Runtime：模型交换的中间桥梁。你训练好的PyTorch模型，转成ONNX，就能在多种推理引擎上跑。省心。
  • TensorRT：NVIDIA家的推理加速引擎。说白了，就是让你的模型在GPU上跑得更快。我实测过，有些模型能提速3-5倍。

核心思路：用Docker把环境打包好，避免“在我电脑上能跑”的尴尬。然后用ONNX做模型转换，最后用TensorRT做推理加速。这套流程，我在多个嵌入式项目里验证过，稳定可靠。

2.2 安装Ubuntu（如果你还没有）

我个人习惯用Ubuntu 20.04 LTS。为什么？因为TensorRT 8.x对它的支持最好。你非要装22.04也行，但有些库的兼容性需要自己折腾。

安装步骤我就不啰嗦了，网上教程一大堆。我只提醒几个关键点：

  • 磁盘分区：建议给根目录至少50GB。我一开始只分了30GB，结果装了几个库就满了，后来还得重新分区，麻烦。
  • 显卡驱动：装完系统第一件事，装NVIDIA驱动。别用系统自带的nouveau，性能差得离谱。
  • 换源：把apt源换成国内的，比如清华源或阿里源。不然下载速度会让你怀疑人生。

小技巧：装完驱动后，用 nvidia-smi 命令检查一下。如果能看到显卡信息，说明驱动装好了。看不到？嗯，重启一下试试。

2.3 安装Docker与NVIDIA Container Toolkit

Docker的安装很简单，一行命令搞定：

curl -fsSL https://get.docker.com -o get-docker.sh
sudo sh get-docker.sh

装完之后，记得把当前用户加到docker组里，不然每次都要sudo：

sudo usermod -aG docker $USER
newgrp docker

接下来是关键——安装NVIDIA Container Toolkit。这个工具能让Docker容器里调用宿主机的GPU。没有它，你在容器里跑TensorRT就是做梦。

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

sudo apt-get update
sudo apt-get install -y nvidia-docker2
sudo systemctl restart docker

验证一下：

docker run --rm --gpus all nvidia/cuda:11.8-base nvidia-smi

如果能看到显卡信息，恭喜你，环境通了。

注意：我曾经遇到过一个问题——容器里跑 nvidia-smi 报错说找不到设备。后来发现是NVIDIA驱动版本太老，升级到525版本后就正常了。如果你也遇到类似问题，先检查驱动版本。

2.4 拉取并配置ONNX Runtime Docker镜像

ONNX Runtime官方提供了Docker镜像，省去了我们自己编译的麻烦。我个人推荐用 onnxruntime:latest-cuda 这个版本，它已经集成了CUDA支持。

docker pull mcr.microsoft.com/onnxruntime:latest-cuda

拉取完成后，启动容器：

docker run -it --rm --gpus all \
  -v /path/to/your/project:/workspace \
  mcr.microsoft.com/onnxruntime:latest-cuda \
  /bin/bash

这里 -v 参数是把宿主机的项目目录挂载到容器里。这样你在宿主机上改代码，容器里直接就能用，不用来回拷贝文件。

验证ONNX Runtime是否正常工作：

python3 -c "import onnxruntime; print(onnxruntime.get_device())"

如果输出 GPU，说明一切正常。如果输出 CPU，嗯，检查一下是不是忘了加 --gpus all。

2.5 安装TensorRT

TensorRT的安装稍微麻烦一点。因为它不是开源的，需要从NVIDIA官网下载。不过别担心，我教你一个省事的办法——直接用NVIDIA官方提供的TensorRT Docker镜像。

docker pull nvcr.io/nvidia/tensorrt:22.12-py3

这个镜像里已经包含了TensorRT、CUDA、cuDNN，以及Python的TensorRT绑定。开箱即用。

启动容器：

docker run -it --rm --gpus all \
  -v /path/to/your/project:/workspace \
  nvcr.io/nvidia/tensorrt:22.12-py3 \
  /bin/bash

验证TensorRT版本：

python3 -c "import tensorrt as trt; print(trt.__version__)"

如果能看到版本号，比如 8.5.3，说明安装成功。

避坑指南：我曾经在TensorRT 8.4上踩过一个坑——某些算子不支持，导致模型转换失败。后来升级到8.5就解决了。所以我的建议是：尽量用最新稳定版，别用太老的版本。

2.6 环境验证：跑一个简单的语义分割模型

环境搭好了，咱们得验证一下能不能用。我准备了一个轻量级语义分割模型——ESNet（Efficient Semantic Network），它只有1.2M参数，非常适合嵌入式场景。

首先，把PyTorch模型转成ONNX：

import torch
import torch.onnx

# 假设你有一个训练好的ESNet模型
model = ESNet(num_classes=19)
model.load_state_dict(torch.load('esnet.pth'))
model.eval()

# 构造一个输入张量
dummy_input = torch.randn(1, 3, 512, 1024)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    'esnet.onnx',
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size', 2: 'height', 3: 'width'},
                  'output': {0: 'batch_size', 2: 'height', 3: 'width'}}
)
print("ONNX模型导出成功！")

然后，用ONNX Runtime验证一下：

import onnxruntime
import numpy as np

# 创建ONNX Runtime会话
ort_session = onnxruntime.InferenceSession('esnet.onnx')

# 构造输入
input_data = np.random.randn(1, 3, 512, 1024).astype(np.float32)

# 推理
outputs = ort_session.run(None, {'input': input_data})
print(f"ONNX Runtime推理完成，输出形状：{outputs[0].shape}")

最后，用TensorRT做加速推理：

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

# 创建TensorRT构建器
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# 创建网络定义
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 加载ONNX模型
parser = trt.OnnxParser(network, logger)
with open('esnet.onnx', 'rb') as f:
    parser.parse(f.read())

# 构建引擎
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
engine = builder.build_serialized_network(network, config)

# 保存引擎
with open('esnet.trt', 'wb') as f:
    f.write(engine)

print("TensorRT引擎构建成功！")

验证结果：如果以上三步都顺利跑通，说明你的环境已经完全就绪。接下来就可以开始真正的模型部署工作了。

2.7 常见问题与解决方案

环境搭建过程中，我遇到过不少问题。这里列几个最常见的：

问题 原因 解决方案
Docker容器里无法调用GPU NVIDIA Container Toolkit未安装或驱动版本不匹配 重新安装nvidia-docker2，并检查驱动版本
ONNX Runtime报错：找不到CUDA ONNX Runtime版本与CUDA版本不兼容 使用官方推荐的版本组合，比如ONNX Runtime 1.14 + CUDA 11.8
TensorRT构建引擎时内存不足 Workspace设置太小 增大 config.set_memory_pool_limit 的值，或者减少batch size
模型转换后精度下降 FP16或INT8量化导致精度损失 使用FP32模式，或者对量化后的模型做校准

特别提醒：如果你在Windows上用WSL2，记得在WSL2里也安装NVIDIA驱动。我有个朋友折腾了两天，最后发现是WSL2的驱动没装。嗯，这种低级错误，咱们就别犯了。

2.8 本章小结

这一章，咱们把Ubuntu、Docker、ONNX Runtime、TensorRT这套环境完整地搭起来了。说实话，环境搭建是嵌入式部署里最枯燥的一步，但也是最重要的一步。环境稳了，后面的工作才能顺利推进。

下一章，我会带你深入语义分割模型的结构，看看怎么把一个PyTorch模型一步步优化到能在嵌入式设备上跑起来。到时候，咱们会用到今天搭好的这套环境。

对了，如果你在搭建过程中遇到任何问题，欢迎随时交流。毕竟，这些坑我都踩过，能帮你省点时间也是好的。

← 上一章 📖 返回目录 下一章 →