NanoTrack原理与转tensorrt推理

最新推荐文章于 2025-07-14 21:02:21 发布

原创最新推荐文章于 2025-07-14 21:02:21 发布 · 置顶 · 4.2k 阅读

46 ·

CC 4.0 BY-SA版权

文章标签：

#1024程序员节

Jetson 同时被 2 个专栏收录

6 篇文章

订阅专栏

目标跟踪

5 篇文章

订阅专栏

部署运行你感兴趣的模型镜像

在这里插入图片描述

文章目录

前言
一、NanoTrack 工作原理
二、运行demo与转换tensorrt模型
三、推理速度对比
总结

前言

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

一、NanoTrack 工作原理

NanoTrackV1 Architecture
示例：pandas 是基于NumPy 的一种工具，该工具是为了解决数据分析任务而创建的。

模型架构:NanoTrack 使用两个相同的卷积神经网络分别处理模板图像（target template）和搜索区域图像（search region）。这两个网络提取的特征随后用于相似性计算。
模板分支与搜索分支:
模板分支：处理目标的模板图像，通常是目标的初始帧或之前帧的一个小区域。通过卷积操作，提取模板的深度特征表示。模板T的input输入为[1,3,127,127]，经过backbone特征提取网络得到output1[1,48,8,8]的模板特征；
搜索分支：处理搜索区域图像，该图像包含当前帧中可能包含目标的较大区域。这个分支也通过卷积操作提取深度特征。图像S的input输入为[1,3,255,255]，经过backbone特征提取网络得到output2[1,48,16,16]的图像特征；
相关性计算:特征图上的每个位置都进行一次简单的元素级相乘，通过这种方式生成一个新的响应图。这个响应图每个像素的位置反映了该位置上的模板和搜索区域特征的相似度,表示目标在搜索区域中的可能位置。
边界框回归:在响应图生成后，跟踪器会根据该响应图来估计目标的位置和尺寸。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-xf
 backbone_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-xf
 backbone_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.py
 head_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部分也要修改：

完成以上运行：