LightGlueProject/LightGlue_Deployment/TENSORRT_QUICKSTART.md

TensorRT 优化快速开始指南

概述

TensorRT 是 NVIDIA 的深度学习推理优化库，可以将 PyTorch 模型优化 2-5 倍速度。

方案选择

方案 1：Torch-TensorRT（推荐，最简单）

优点：

• 安装简单
• 与 PyTorch 无缝集成
• 自动优化

缺点：

• 可能不支持所有操作
• 首次编译较慢

方案 2：ONNX → TensorRT（更稳定）

优点：

• 更稳定
• 支持更多模型

缺点：

• 需要两步转换（ONNX → TensorRT）
• 配置更复杂

快速开始（方案 1：Torch-TensorRT）

步骤 1：安装依赖

# 安装 Torch-TensorRT
pip install torch-tensorrt

# 验证安装
python -c "import torch_tensorrt; print('TensorRT installed successfully')"

步骤 2：修改代码

在 demo_lightglue_camera_position_async.py 中添加：

1. 在文件开头添加导入：

try:
    from tensorrt_wrapper import create_tensorrt_models
    TENSORRT_AVAILABLE = True
except ImportError:
    TENSORRT_AVAILABLE = False

2. 在 argparse 中添加参数：

parser.add_argument("--use_tensorrt", action="store_true", 
                   help="Use TensorRT optimized models")
parser.add_argument("--tensorrt_precision", type=str, default="fp16",
                   choices=["fp32", "fp16", "int8"],
                   help="TensorRT precision mode")

3. 在模型加载后（约第338行）添加：

print("Loaded SuperPoint and LightGlue models")

# TensorRT 优化
if opt.use_tensorrt and TENSORRT_AVAILABLE and device == "cuda":
    try:
        print("Compiling models with TensorRT...")
        extractor, matcher = create_tensorrt_models(
            extractor, matcher, precision=opt.tensorrt_precision
        )
        print("✓ TensorRT models compiled successfully")
    except Exception as e:
        print(f"✗ Failed to compile with TensorRT: {e}")
        print("Falling back to PyTorch models")

步骤 3：运行程序

# 首次运行（会编译模型，需要5-15分钟）
python demo_lightglue_camera_position_async.py \
    --input "udp://0.0.0.0:12346" \
    --max_keypoints 128 \
    --use_fp16 \
    --use_tensorrt \
    --tensorrt_precision fp16

# 第二次运行（直接加载编译好的模型，很快）
python demo_lightglue_camera_position_async.py \
    --input "udp://0.0.0.0:12346" \
    --max_keypoints 128 \
    --use_tensorrt

性能预期

模式	FPS	提升
PyTorch FP32	~15-18	基准
PyTorch FP16	~22-24	+30%
TensorRT FP16	~35-45	+100-150%
TensorRT INT8	~50-60	+200-250%

故障排除

问题 1：编译失败

错误： torch_tensorrt not found

解决：

pip install torch-tensorrt

问题 2：CUDA 错误

错误： CUDA out of memory 或 CUDA version mismatch

解决：

• 检查 CUDA 版本：nvidia-smi
• 确保 PyTorch 和 TensorRT 使用相同的 CUDA 版本
• 减少 workspace_size 或 max_keypoints

问题 3：模型不支持的操作

错误： Unsupported operation: XXX

解决：

• 某些操作可能不被 TensorRT 支持
• 尝试使用 torch_executed_ops 参数排除这些操作
• 或使用 ONNX → TensorRT 方案

问题 4：精度下降

解决：

• 使用 FP16 而不是 INT8
• 检查模型输出是否正常
• 可能需要重新训练或微调

高级选项

自定义编译选项

在 tensorrt_wrapper.py 中修改：

self.trt_model = torch_tensorrt.compile(
    self.model,
    inputs=example_input,
    enabled_precisions={torch.half},
    workspace_size=2 << 30,  # 2GB（增加内存）
    min_block_size=7,  # 最小块大小
    torch_executed_ops=["unsupported_op"],  # 排除不支持的操作
)

使用 INT8 量化

# 需要校准数据（使用真实图像）
python demo_lightglue_camera_position_async.py \
    --use_tensorrt \
    --tensorrt_precision int8

文件说明

• tensorrt_wrapper.py: TensorRT 包装器
• convert_to_tensorrt.py: ONNX → TensorRT 转换脚本
• tensorrt_optimization_guide.md: 详细指南
• tensorrt_integration_example.py: 集成示例

下一步

1. 测试 TensorRT FP16 性能
2. 如果还不够快，尝试 INT8 量化
3. 如果 Torch-TensorRT 失败，尝试 ONNX → TensorRT 方案

参考

• Torch-TensorRT 文档
• TensorRT 文档
• LightGlue 项目