3. 图像分类模型部署

任务目标

在Jetson Nano上使用TensorRT部署预训练的ResNet-18模型，对单张静态图片执行推理，输出Top-5类别名称及对应置信度。本案例聚焦于从PyTorch模型到TensorRT引擎的完整转换与推理闭环，不涉及训练过程。

环境与依赖

  • Jetson Nano（JetPack 4.6+，已预装TensorRT 8.x、CUDA 11.x）
  • PyTorch 1.10+（用于导出ONNX）
  • Python 3.6+，依赖：torch, torchvision, onnx, pycuda, Pillow, numpy
  • ImageNet类别映射文件（imagenet_classes.txt，可从下方代码自动下载）

实现代码

以下代码完成从PyTorch预训练模型导出ONNX、构建TensorRT FP16引擎、预处理输入图片并输出Top-5结果的全部流程。将代码保存为trt_resnet18.py。

import torch
import torchvision.models as models
import torchvision.transforms as transforms
import onnx
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
from PIL import Image
import urllib.request
import os
import time

# ---------- 配置 ----------
ONNX_PATH = "resnet18.onnx"
ENGINE_PATH = "resnet18_fp16.engine"
IMAGE_PATH = "test.jpg"  # 输入图片
NUM_CLASSES = 1000
TOP_K = 5
INPUT_SIZE = 224

# ---------- 1. 导出ONNX ----------
def export_onnx():
    model = models.resnet18(pretrained=True).eval().cuda()
    dummy = torch.randn(1, 3, INPUT_SIZE, INPUT_SIZE).cuda()
    torch.onnx.export(model, dummy, ONNX_PATH,
                      input_names=["input"], output_names=["output"],
                      opset_version=11, dynamic_axes=None)
    print(f"ONNX已导出: {ONNX_PATH}")

# ---------- 2. 构建TensorRT引擎 ----------
def build_engine():
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, 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)
    input_tensor.shape = [1, 3, INPUT_SIZE, INPUT_SIZE]
    profile.set_shape(input_tensor.name, [1, 3, INPUT_SIZE, INPUT_SIZE], 
                      [1, 3, INPUT_SIZE, INPUT_SIZE], [1, 3, INPUT_SIZE, INPUT_SIZE])
    config.add_optimization_profile(profile)
    
    engine = builder.build_serialized_network(network, config)
    with open(ENGINE_PATH, "wb") as f:
        f.write(engine)
    print(f"TensorRT引擎已保存: {ENGINE_PATH}")

# ---------- 3. 加载ImageNet类别 ----------
def load_labels():
    if not os.path.exists("imagenet_classes.txt"):
        url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
        urllib.request.urlretrieve(url, "imagenet_classes.txt")
    with open("imagenet_classes.txt") as f:
        return [line.strip() for line in f.readlines()]

# ---------- 4. 图片预处理 ----------
def preprocess(image_path):
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(INPUT_SIZE),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    img = Image.open(image_path).convert("RGB")
    tensor = transform(img).unsqueeze(0).numpy()
    return np.ascontiguousarray(tensor, dtype=np.float32)

# ---------- 5. TensorRT推理 ----------
def infer():
    # 加载引擎
    logger = trt.Logger(trt.Logger.WARNING)
    runtime = trt.Runtime(logger)
    with open(ENGINE_PATH, "rb") as f:
        engine = runtime.deserialize_cuda_engine(f.read())
    context = engine.create_execution_context()
    
    # 分配内存
    input_shape = (1, 3, INPUT_SIZE, INPUT_SIZE)
    output_shape = (1, NUM_CLASSES)
    d_input = cuda.mem_alloc(trt.volume(input_shape) * trt.float32.itemsize)
    d_output = cuda.mem_alloc(trt.volume(output_shape) * trt.float32.itemsize)
    stream = cuda.Stream()
    
    # 预处理输入
    input_data = preprocess(IMAGE_PATH)
    output_data = np.empty(output_shape, dtype=np.float32)
    
    # 拷贝到设备并推理
    cuda.memcpy_htod_async(d_input, input_data, stream)
    context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(output_data, d_output, stream)
    stream.synchronize()
    
    # 后处理Top-5
    labels = load_labels()
    probs = torch.from_numpy(output_data).softmax(dim=1)[0]
    top5_prob, top5_idx = torch.topk(probs, TOP_K)
    
    print(f"\n===== Top-{TOP_K} 分类结果 =====")
    for i in range(TOP_K):
        idx = top5_idx[i].item()
        print(f"{i+1}. {labels[idx]:<30} 置信度: {top5_prob[i].item():.4f}")

if __name__ == "__main__":
    if not os.path.exists(ONNX_PATH):
        export_onnx()
    if not os.path.exists(ENGINE_PATH):
        build_engine()
    infer()

运行与验证

  1. 准备测试图片：将任意一张224x224以上的JPEG图片命名为test.jpg，放在脚本同目录。
  2. 执行推理：python3 trt_resnet18.py
  3. 首次运行会自动下载ImageNet类别文件（约30KB），并依次完成ONNX导出与TensorRT引擎构建（FP16，耗时约1-2分钟）。
  4. 再次运行会直接加载已构建的引擎，跳过导出与构建阶段。

通过判定标准：程序输出5行结果，每行包含类别名称（如“golden retriever”）和0到1之间的置信度值，且Top-1置信度通常高于0.5（若图片内容明确）。若图片内容模糊或类别不在ImageNet 1000类中，置信度可能偏低，但输出格式正确且无报错即视为部署成功。

常见问题

  • 显存不足：构建引擎时若报CUDA out of memory，将1 << 30（1GB）改为1 << 29（512MB）。
  • ONNX解析失败：确认PyTorch版本≥1.10，且opset_version=11兼容。
  • 推理速度慢：确认引擎文件已生成（非重新构建），且Jetson Nano处于最高性能模式（sudo nvpmodel -m 0）。
← 上一章 📖 返回目录 下一章 →