NanoTrack原理与转tensorrt推理

---

前言

NanoTrack 是一种轻量级且高效的目标跟踪算法，基于Siamese网络架构，旨在在资源受限的设备（如无人机、嵌入式系统等）上实现实时且精确的目标跟踪。随着计算机视觉领域的快速发展，传统的目标跟踪算法（如KCF等）逐渐暴露出在复杂场景中的局限性，特别是在处理目标的尺度变化、遮挡和背景复杂时，性能难以保证。
本文档将详细介绍 NanoTrack 的模型结构以帮助读者更好地理解并应用该算法，同时对比pt、onnx、tensorrt进行推理的耗时情况。
NanoTrack源码下载：https://github.com/HonglinChu/SiamTrackers/tree/master/NanoTrack

---

一、NanoTrack 工作原理

NanoTrackV1 Architecture

1. 模型架构:NanoTrack 使用两个相同的卷积神经网络分别处理模板图像（target template）和搜索区域图像（search region）。这两个网络提取的特征随后用于相似性计算。

2. 模板分支与搜索分支:
   模板分支：处理目标的模板图像，通常是目标的初始帧或之前帧的一个小区域。通过卷积操作，提取模板的深度特征表示。模板T的input输入为[1,3,127,127]，经过backbone特征提取网络得到output1[1,48,8,8]的模板特征；
   搜索分支：处理搜索区域图像，该图像包含当前帧中可能包含目标的较大区域。这个分支也通过卷积操作提取深度特征。图像S的input输入为[1,3,255,255]，经过backbone特征提取网络得到output2[1,48,16,16]的图像特征；

3. 相关性计算:特征图上的每个位置都进行一次简单的元素级相乘，通过这种方式生成一个新的响应图。这个响应图每个像素的位置反映了该位置上的模板和搜索区域特征的相似度,表示目标在搜索区域中的可能位置。

4. 边界框回归:在响应图生成后，跟踪器会根据该响应图来估计目标的位置和尺寸。Pointwise-Correlation 的输出会传递给回归网络，预测目标的最终边界框。

二、运行demo与转换tensorrt模型

2.1 运行pt模型demo

配置环境这些基础的步骤这里就不做赘述了。
以下讲解以V3源码为例：
在NanoTrack/bin/demo.py里修改配置文件与模型路径，改成V3版本的：

parser = argparse.ArgumentParser(description='tracking demo') parser.add_argument('--config', default='../models/config/configv3.yaml',type=str, help='config file')parser.add_argument('--snapshot', default='../models/pretrained/nanotrackv3.pth', type=str, help='model name')parser.add_argument('--video_name', default='../bin/girl_dance.mp4', type=str, help='videos or image files')parser.add_argument('--save', action='store_true', help='whether visualzie result') args = parser.parse_args()

然后运行：

python demo.py

2.2 转onnx模型

这里将pt模型拆成3个部分，在NanoTrack/pytorch2onnx.py中进行修改：

# backbone input-xfbackbone_s = torch.randn([1, 3, 255, 255], device=device)export_onnx_file_path = './models/onnx/nanotrack_backbone_S.onnx'torch.onnx.export(backbone_net, backbone_s, export_onnx_file_path, input_names=['input'], output_names=['output'],verbose=True, opset_version=14)# backbone input-xfbackbone_x = torch.randn([1, 3, 127, 127], device=device)export_onnx_file_path = './models/onnx/nanotrack_backbone_T.onnx'torch.onnx.export(backbone_net, backbone_x, export_onnx_file_path, input_names=['input'], output_names=['output'],verbose=True, opset_version=14)# head  change forward  /media/dell/Data/NanoTrack/nanotrack/models/model_builder.pyhead_zf, head_xf = torch.randn([1, 96, 8, 8], device=device), torch.randn([1, 96, 16, 16], device=device)export_onnx_file_path = './models/onnx/nanotrack_head.onnx'torch.onnx.export(head_net, (head_zf, head_xf), export_onnx_file_path, input_names=['input1', 'input2'],output_names=['output1', 'output2'], verbose=True, opset_version=14)

这里需要注意V3模型，模板T的input输入为[1,3,127,127]，经过backbone特征提取网络得到output1[1,96,8,8]的模板特征；

图像S的input输入为[1,3,255,255]，经过backbone特征提取网络得到output2[1,96,16,16]的图像特征；

output1和output2会作为input1与input2被送入模型head部分，v3默认输出是output1[1,2,15,15],output2[1,4,15,15]

具体的可以查询NanoTrack/models/config/configv3.yaml中的：

2.3 转tensorrt模型

下载NanoTrack TRT项目:https://github.com/SSSSSSL/NanoTrack_TensorRT/tree/master

这里根据自己转出的onnx模型为例，进行修改路径（也可以直接运行NanoTrack_TensorRT/create_trt_engine.sh生成项目自带的模型）

BACK_EXAM_ONNX="models/nanotrack_backbone_S.onnx"
BACK_TEMP_ONNX="models/nanotrack_backbone_T.onnx"
HEAD_ONNX="models/nanotrack_head.onnx"TRT_BIN="/usr/src/tensorrt/bin/trtexec"BACK_EXAM_TRT="models/nanotrack_backbone_S.engine"
BACK_TEMP_TRT="models/nanotrack_backbone_T.engine"
HEAD_TRT="models/nanotrack_head.engine"

运行NanoTrack_TensorRT/create_trt_engine.sh

出现上图说明转模型成功

2.4 运行trt模型推理

修改NanoTrack_TensorRT/main.py中的engine文件路径：

def main():""" 主要功能是加载TensorRT引擎并进行推理 """back_exam_engine_path = "models/nanotrack_backbone_S.engine"back_temp_engine_path = "models/nanotrack_backbone_T.engine"head_engine_path = "models/nanotrack_head.engine"tracker = Tracker(back_exam_engine_path, back_temp_engine_path, head_engine_path)# webcam 开始追踪# trackLive(tracker)# video 开始追踪video_path = "/home/nvidia/Desktop/project/video/girl_dance.mp4"trackVideo(tracker, video_path, save_video=True)# benchmark 开始# vot_path = "/home/rgblab/tracking/dataset/VOT"# benchmarkVOT(tracker, vot_path)

因为我们输出的onnx响应图为15*15的，需要在NanoTrack_TensorRT/core/config.py进行相应修改：

相应的NanoTrack_TensorRT/core/process.py中的HeadProcessor的post部分也要修改：

完成以上运行：

python main.py

三、推理速度对比

模型	平均FPS
pt	70
onnx	95
tensorrt	132

---

总结

本文介绍了 NanoTrack 的工作原理，并通过实践展示了如何运行pt模型、转换为ONNX模型以及TensorRT模型，最终在嵌入式设备上进行推理。通过对比不同推理框架的耗时情况，展示了TensorRT模型在加速推理方面的优势（相比于pt模型，TensorRT模型平均FPS提升了约90%）。

参考文档：
https://github.com/HonglinChu/SiamTrackers/tree/master/NanoTrack
https://github.com/SSSSSSL/NanoTrack_TensorRT/tree/master
https://blog.csdn.net/X_windy/article/details/126153227
https://blog.csdn.net/qq_58248345/article/details/135825911

如果阅读本文对你有用，欢迎一键三连呀！！！
2024年10月24日08:48:03