目标跟踪之ICCV2021 Learning Spatio-Temporal Transformer for Visual Tracking代码复现

代码地址:https://github.com/researchmm/Stark
论文地址:https://arxiv.org/abs/2103.17154

一、配置环境

conda create -n stark python=3.6
conda activate stark
bash install_pytorch17.sh

二、设置工程路径

  1. 运行指令python tracking/create_default_local_file.py --workspace_dir . --data_dir ./data --save_dir .
  2. 修改lib/test/evaluation/local.py文件,设置数据集路径
  3. 新建文件夹checkpoints,权重文件保存地址如下:
${STARK_ROOT}
 -- checkpoints
     -- train
         -- stark_s
             -- baseline
                 STARKS_ep0500.pth.tar
             -- baseline_got10k_only
                 STARKS_ep0500.pth.tar
         -- stark_st2
             -- baseline
                 STARKST_ep0050.pth.tar
             -- baseline_got10k_only
                 STARKST_ep0050.pth.tar
             -- baseline_R101
                 STARKST_ep0050.pth.tar
             -- baseline_R101_got10k_only
                 STARKST_ep0050.pth.tar

三、数据集评测

作者一共给了两个版本,四个模型:
百度云下载链接:https://pan.baidu.com/s/1gIbGAejusJ_aaD82Pd1rrQ
提取码:0krk
运行stark_st2/baseline/STARKST_ep0050.pth.tar对应的是STARK-ST50,运行指令python tracking/test.py stark_st baseline --dataset trackingnet --threads 0

在这里插入图片描述

另一个版本是stark_st2/baseline_R101/STARKST_ep0050.pth.tar,对应论文的STARK-ST101,运行指令python tracking/test.py stark_st baseline_R101 --dataset trackingnet --threads 0
在这里插入图片描述
STARK-ST101速度还行,1080Ti上13帧/s。
STARK-ST50速度稍微快一点,18帧/s。

四、代码训练

修改数据集路径:lib/train/admin/local.py
作者默认的四个训练集是['LASOT', 'GOT10K_vottrain', 'COCO17', 'TRACKINGNET']
如果要修改数据集,修改cfg文件'Stark-main/experiments/stark_st1/baseline_R101.yaml'进行修改:
目标跟踪之ICCV2021 Learning Spatio-Temporal Transformer for Visual Tracking代码复现_第1张图片
Stark的训练主要有两个阶段,根据原文的描述,第一阶段stark_st1主要是用于定位精度(角点预测),第二阶段stark_st2用于提升stark的分类能力(鉴别能力)。
目标跟踪之ICCV2021 Learning Spatio-Temporal Transformer for Visual Tracking代码复现_第2张图片

单卡训练指令: python tracking/train.py --script stark_st1 --config baseline_R101 --save_dir . --mode single
第一阶段训练500epochs,之后再进行第二阶段的50epochs训练python tracking/train.py --script stark_st2 --config baseline_R101 --save_dir . --mode single --script_prv stark_st1 --config_prv baseline_R101
目标跟踪之ICCV2021 Learning Spatio-Temporal Transformer for Visual Tracking代码复现_第3张图片
目标跟踪之ICCV2021 Learning Spatio-Temporal Transformer for Visual Tracking代码复现_第4张图片
第一阶段训练结束如下:
目标跟踪之ICCV2021 Learning Spatio-Temporal Transformer for Visual Tracking代码复现_第5张图片
第二阶段训练:(注意要把测试时,保存在checkpoints/train/stark_st2/baseline_R101/底下的权重文件换一个位置保存,不然会直接读取,导致网络错误,无法训练,正常开启训练如下)
目标跟踪之ICCV2021 Learning Spatio-Temporal Transformer for Visual Tracking代码复现_第6张图片
训练时间特别久,单张1080Ti断断续续感觉得花十天时间。建议有条件的多卡一起训练。

你可能感兴趣的:(目标跟踪,python,深度学习)