📚 课程目录第4章基于TensorRT的GPU边缘设备推理优化

4. 基于TensorRT的GPU边缘设备推理优化

任务目标

在Jetson Nano（JetPack 4.6+）上，将一个PyTorch图像分类模型转换为TensorRT FP16引擎文件，并编写推理脚本完成端到端验证。最终通过对比FP32与FP16的推理耗时与输出一致性，确认优化生效。

环境与依赖

Jetson Nano开发套件（JetPack 4.6，含CUDA 10.2、cuDNN 8、TensorRT 8.2）
PyTorch 1.10（仅用于导出ONNX，转换完成后可卸载）
Python 3.6+，需安装：torch、onnx、onnxruntime-gpu（可选，用于对比验证）、pycuda、numpy、Pillow

实现步骤

整体流程为：PyTorch模型 → ONNX → TensorRT引擎（FP16）→ 推理验证。以下代码在Jetson Nano上直接执行，核心包含两个脚本：build_engine.py（转换）与infer.py（推理验证）。

1. 导出ONNX并构建TensorRT引擎

使用ResNet18作为示例模型，输入尺寸固定为3×224×224。构建引擎时开启FP16模式，并设置工作空间大小。

# build_engine.py
import torch
import torchvision.models as models
import onnx
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

# ---------- 1. PyTorch → ONNX ----------
model = models.resnet18(pretrained=True).eval().cuda()
dummy_input = torch.randn(1, 3, 224, 224).cuda()
onnx_path = "resnet18.onnx"
torch.onnx.export(model, dummy_input, onnx_path,
                  input_names=["input"], output_names=["output"],
                  opset_version=11, dynamic_axes=None)
print("ONNX导出完成")

# ---------- 2. ONNX → TensorRT FP16引擎 ----------
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)

with open(onnx_path, "rb") as f:
    assert parser.parse(f.read()), "ONNX解析失败"

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
config.set_flag(trt.BuilderFlag.FP16)  # 启用FP16

# 固定batch size为1
profile = builder.create_optimization_profile()
input_tensor = network.get_input(0)
profile.set_shape(input_tensor.name, (1, 3, 224, 224), (1, 3, 224, 224), (1, 3, 224, 224))
config.add_optimization_profile(profile)

engine_bytes = builder.build_serialized_network(network, config)
if engine_bytes is None:
    raise RuntimeError("引擎构建失败")

with open("resnet18_fp16.engine", "wb") as f:
    f.write(engine_bytes)
print("FP16引擎已保存: resnet18_fp16.engine")

2. 使用TensorRT引擎进行FP16推理

加载引擎文件，分配GPU显存，执行前向推理。代码同时支持加载ONNX Runtime进行FP32精度对比（可选）。

# infer.py
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import time
from PIL import Image
import torchvision.transforms as transforms

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

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

def allocate_buffers(engine):
    h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=np.float32)
    h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=np.float32)
    d_input = cuda.mem_alloc(h_input.nbytes)
    d_output = cuda.mem_alloc(h_output.nbytes)
    stream = cuda.Stream()
    return h_input, h_output, d_input, d_output, stream

def infer(engine, input_image):
    h_input, h_output, d_input, d_output, stream = allocate_buffers(engine)
    context = engine.create_execution_context()
    h_input[:] = input_image.ravel()
    cuda.memcpy_htod_async(d_input, h_input, stream)
    context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    stream.synchronize()
    return h_output.copy()

# ---------- 图像预处理 ----------
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

img = Image.open("test_image.jpg").convert("RGB")
input_tensor = transform(img).unsqueeze(0).numpy()

# ---------- TensorRT FP16推理 ----------
engine = load_engine("resnet18_fp16.engine")
# 预热
_ = infer(engine, input_tensor)

# 计时推理（100次取平均）
times = []
for _ in range(100):
    t0 = time.time()
    output = infer(engine, input_tensor)
    times.append(time.time() - t0)
trt_time = np.mean(times[10:]) * 1000  # 跳过前10次预热

# 获取Top-1类别
pred_class = int(np.argmax(output))
print(f"TensorRT FP16 平均推理耗时: {trt_time:.2f} ms")
print(f"预测类别索引: {pred_class}")

运行与验证

将上述两个脚本保存至Jetson Nano同一目录，并准备一张测试图片test_image.jpg。
执行转换：python3 build_engine.py，看到“FP16引擎已保存”即成功。
执行推理：python3 infer.py，观察输出耗时与类别索引。

通过条件：

引擎文件成功生成且无报错；
推理脚本输出一个0~999之间的类别索引，且与PyTorch原模型（FP32）在相同输入上的预测结果一致；
FP16推理耗时明显低于FP32（通常可降低30%~50%）。若需定量对比，可在infer.py中额外加载ONNX Runtime FP32模型，用相同输入测量耗时并比较。

常见问题

显存不足：降低WORKSPACE大小至512MB，或减小batch size。
FP16精度损失：若分类结果与FP32不一致，可尝试对输入做归一化检查，或改用INT8（需校准数据集）。
ONNX导出失败：确认PyTorch版本≥1.8，且opset_version≥11。