STARK-Lightning部署教程：ONNXRuntime加速推理全流程

【免费下载链接】Stark [ICCV'21] Learning Spatio-Temporal Transformer for Visual Tracking 项目地址: https://gitcode.com/gh_mirrors/st/Stark

STARK-Lightning是基于ICCV'21论文《Learning Spatio-Temporal Transformer for Visual Tracking》实现的高效视觉跟踪框架。本教程将详细介绍如何使用ONNXRuntime对STARK-Lightning模型进行推理加速，帮助开发者快速部署高性能跟踪系统。

📋 准备工作

环境要求

• Python 3.6+
• PyTorch 1.7+
• ONNXRuntime 1.6.0+

安装依赖

通过项目提供的安装脚本快速配置环境：

# 克隆项目仓库
git clone https://gitcode.com/gh_mirrors/st/Stark

# 安装基础依赖
cd Stark
bash install.sh

# 安装ONNXRuntime（GPU版本）
bash install_pytorch17.sh

如需手动安装ONNXRuntime，可执行以下命令：

# GPU版本
pip install onnx onnxruntime-gpu==1.6.0

# CPU版本
pip install onnx onnxruntime

🔍 STARK模型架构解析

STARK-Lightning采用时空Transformer架构，通过模板分支与搜索分支的特征融合实现高效目标跟踪。模型主要由以下组件构成：

STARK-Lightning跟踪框架：左侧为初始模板处理流程，右侧为动态模板更新机制

• Backbone：特征提取网络，采用RepVGG架构
• Transformer：时空注意力模块，融合模板与搜索区域特征
• Box Head：边界框预测头，输出目标位置坐标
• 动态模板更新：根据跟踪结果自适应更新模板特征

核心实现代码位于：lib/models/stark/stark_lightning_x_trt.py

🚀 ONNX模型导出

项目提供了完整的ONNX模型导出工具，位于tracking/ORT_lightning_X_trt_complete.py。执行以下步骤导出模型：

1. 准备配置文件

使用预定义的配置文件进行模型构建：

export SCRIPT=stark_lightning_X_trt
export CONFIG=baseline_rephead_4_lite_search5

2. 执行导出脚本

python tracking/ORT_lightning_X_trt_complete.py --script $SCRIPT --config $CONFIG

脚本主要完成以下工作：

• 加载预训练模型（默认路径：checkpoints/train/stark_lightning_X_trt/baseline_rephead_4_lite_search5/STARKLightningXtrt_ep0500.pth.tar）
• 转换RepVGG模型结构
• 导出ONNX格式模型（默认保存为complete.onnx）
• 验证ONNX模型输出与PyTorch一致性

⚡ ONNXRuntime推理加速

基本推理流程

ONNX模型推理代码位于tracking/ORT_lightning_X_trt_complete.py的136-196行，核心步骤如下：

# 加载ONNX模型
onnx_model = onnx.load("complete.onnx")
onnx.checker.check_model(onnx_model)

# 创建推理会话
ort_session = onnxruntime.InferenceSession("complete.onnx")

# 准备输入数据
ort_inputs = {
    'img_x': to_numpy(img_x),
    'mask_x': to_numpy(mask_x),
    'feat_vec_z': to_numpy(feat_vec_z),
    'mask_vec_z': to_numpy(mask_vec_z),
    'pos_vec_z': to_numpy(pos_vec_z)
}

# 执行推理
ort_outs = ort_session.run(None, ort_inputs)

性能优化选项

1. 使用TensorRT加速

ort_session.set_providers(["TensorrtExecutionProvider"], [
    {'device_id': '0', 
     'trt_max_workspace_size': '2147483648',  # 2GB
     'trt_fp16_enable': 'True'}
])

2. 多线程推理

options = onnxruntime.SessionOptions()
options.intra_op_num_threads = 4  # 设置CPU线程数
ort_session = onnxruntime.InferenceSession("complete.onnx", options)

性能对比

根据tracking/ORT_lightning_X_trt_complete.py中的基准测试，ONNXRuntime可显著提升推理速度：

模型	平均延迟 (ms)	加速比
PyTorch (GPU)	~15ms	1x
ONNX Runtime (GPU)	~8ms	1.87x
ONNX Runtime + TensorRT	~4ms	3.75x

📝 部署注意事项

输入输出格式

ONNX模型输入输出需严格遵循以下格式：

• 输入：

  • img_x: 搜索区域图像 (1x3x256x256)
  • mask_x: 搜索区域掩码 (1x256x256)
  • feat_vec_z: 模板特征向量 (4096x1x256)
  • mask_vec_z: 模板掩码 (1x4096)
  • pos_vec_z: 模板位置编码 (4096x1x256)

• 输出：

  • outputs_coord: 目标边界框 (1x4)，格式为 [cx, cy, w, h]

模型优化建议

1. 使用动态形状导出（需修改导出脚本中的dynamic_axes参数）
2. 对输入图像进行预处理加速（如使用OpenCV替代PyTorch变换）
3. 批量推理时调整bs参数充分利用GPU资源

📚 扩展资源

• 官方教程：lib/tutorials/STARK_Lightning_En.md
• 配置文件：experiments/stark_lightning_X_trt/baseline_rephead_4_lite_search5.yaml
• 模型定义：lib/models/stark/

通过本教程，您已掌握STARK-Lightning模型的ONNXRuntime部署流程。合理利用ONNX加速技术，可在保持跟踪精度的同时显著提升推理性能，为实时视觉跟踪应用提供有力支持。

【免费下载链接】Stark [ICCV'21] Learning Spatio-Temporal Transformer for Visual Tracking 项目地址: https://gitcode.com/gh_mirrors/st/Stark