手头有一个新项目要复现一看是ICCV 2021上的跟踪算法Stark全称Learning Spatio-Temporal Transformer for Visual Tracking。这几年Transformer结构在检测、分割里大杀四方跟踪领域也终于等来了这波红利。Stark算是把Transformer架构完整引入单目标跟踪的标志性工作之一在LaSOT、TrackingNet这类长时数据集上的精度相比之前的Siamese类方法有比较明显的提升。这篇就来把它的环境配置、数据准备、训练测试整个流程走一遍重点聊聊我踩过的坑和那些代码里不会写明白的细节。先说结论Stark的代码结构不算复杂但配置门槛集中在三块——环境依赖版本、数据集目录格式、预训练模型的下载与加载。这三块只要有一处没对齐后面训练和测试就会各种报错。下面直接进入正题。1. 配置前的核心认知1.1 算法核心思路速览Stark全称是Learning Spatio-Temporal Transformer for Visual Tracking它跟传统的Siamese跟踪器最大区别在于不再用互相关操作来融合模板和搜索区域而是把模板和搜索区域都当成一组token序列喂给Transformer encoder做全局建模再用一个decoder直接输出目标的包围框。整个网络可以拆成三部分Backbone默认是ResNet-50负责提特征。也有基于Swin Transformer的加强版本但复现最稳的还是ResNet-50这套。Encoder把模板特征和搜索区域特征展平成token序列拼在一起输入Transformer encoder让目标模板和搜索区域在全局范围内做特征交互。Decoder借鉴了DETR的思路用一组可学习的对象查询object query来解码目标位置。Stark的特别之处在于提供了一个IoU预测分支直接回归预测框和真实框的IoU得分推理时用这个得分作为置信度参考。另外动态模板更新机制也是Stark的亮点。简单说它不只用第一帧的模板还会根据历史帧的跟踪结果构造新的时空模板让网络在目标表观变化时依然能稳定跟上。1.2 为什么配置这个词值得单独写一篇很多人觉得跟踪算法的配置无非就是装依赖、改路径、跑训练没必要大动干戈。但Stark跟一般的分类检测代码有一定区别第一它的训练依赖四个数据集每个数据集的目录组织方式都不一样。LaSOT、GOT-10k、TrackingNet、COCO这四者的标注格式和文件夹层级各不相同代码里对不同数据集走了不同的读取分支。第二它的预训练backbone加载逻辑比较特别官方代码中提供了从MAE等预训练方法转化来的权重文件如果你手头改过模型结构加载key会对不上这时需要手动修改代码里的加载逻辑。第三训练和测试分离测试时还需要额外的评估脚本来计算成功率等指标。整个链路涉及多个模块任何一环出问题都会导致结果文件生成失败。2. 环境与依赖准备2.1 硬件与系统要求先看硬件。Stark的backbone是ResNet-50Transformer层数不算多训练时模板尺寸是128×128搜索区域320×320batch size按默认配置是64多卡情况下每卡16。这个显存压力主要来自Transformer encoder对长序列的自注意力计算所以单卡训练建议至少11GB显存也就是GTX 1080Ti或RTX 2080Ti这个级别再往下容易OOM。如果只有8GB显存需要把batch size降到8或4同时可以把搜索区域从320降到288但精度会有一点损失。推理阶段要求低很多6GB显存也能跑因为推理时batch size为1。系统方面Ubuntu 18.04或20.04最稳。Windows不是不行但Cython编译、数据路径处理、多进程数据加载这些环节容易出幺蛾子不建议新手用Windows折腾。2.2 Python环境与依赖版本组合官方推荐基于Python 3.7或3.8PyTorch 1.5以上即可但综合社区反馈和我的实际测试比较省心的组合是Python 3.8.10 PyTorch 1.8.0 CUDA 11.1 torchvision 0.9.0为什么推荐这个组合因为Stark代码里用到的很多API在PyTorch 1.8里都有稳定实现而且官方发布的预训练模型是基于PyTorch 1.7/1.8训练的权重版本跨度太大会在加载时碰到序列化不兼容的警告。PyTorch 2.x也能跑但需要额外注意一些算子兼容性问题没必要自找麻烦。依赖库方面官方requirements.txt里列了timm、yacs、lyft_dataset_sdk等但实际跑通还需要装不少隐式依赖。以下是我整理出的完整依赖清单pip install torch1.8.0 torchvision0.9.0 --extra-index-url https://download.pytorch.org/whl/cu111 pip install numpy opencv-python pandas tqdm pyyaml pip install timm0.4.12 yacs scipy tensorboardX pip install lyft_dataset_sdk attrs pip install pycocotools几个容易出问题的包单独说明一下lyft_dataset_sdk这个包在PyPI上的最新版本可能需要较新的protobuf如果安装时报冲突可以先装protobuf 3.20.x再装它。timm版本很重要。timm 0.4.12和0.5.x之间有部分模型接口改名Stark代码里调用timm的方式是按0.4.x写的装太新的版本会报Model object has no attribute reset_drop_path之类的错误。opencv-python建议4.x不要装opencv-contrib-python有时两者会冲突导致cv2导入失败。2.3 Cython模块编译代码里有一个涉及bbox overlap计算的cython模块需要编译官方代码在lib/utils/cython_bbox下面。编译前需要确认系统里装了gcc和python开发头文件sudo apt-get install build-essential python3-dev cd lib/utils/cython_bbox python setup.py build_ext --inplace我遇到最典型的问题是编译不报错但运行时提示cython_bbox.so文件找不到。这种情况通常是setup.py里输出的so文件命名和代码里import的模块名不一致或者没有正确指定--inplace。解决办法是编译完成后把so文件复制一份改名为cython_bbox.so放到同级目录。3. 数据集准备与目录结构3.1 四个训练数据集的格式要求Stark在训练时使用LaSOT、GOT-10k、TrackingNet和COCO四个数据集。前三个是跟踪数据集COCO用来扩充大规模通用目标样本增强模型泛化能力。这四个数据集的目录结构必须按代码里的预期组织否则训练时读不到数据会报FileNotFoundError或直接卡在数据加载阶段。以官方项目根目录为例推荐的数据集布局是这样的stark/ ├── datasets/ │ ├── lasot/ │ ├── got10k/ │ ├── trackingnet/ │ └── coco/ ├── experiments/ ├── lib/ ├── pretrained_models/ ├── train.py └── tracking/LaSOT的目录结构相对复杂一些lasot/ ├── train/ │ ├── airplane-1/ │ │ ├── img/ │ │ │ ├── 00000001.jpg │ │ │ └── ... │ │ ├── groundtruth.txt │ │ └── full_occlusion.txt │ └── ... ├── testing/ │ └── ...注意LaSOT的groundtruth.txt每行是8个数代表旋转矩形的四个角点坐标而代码内部会把它转换成axis-aligned的框所以不要手动改这个文件保持原样即可。GOT-10k的结构相对简单got10k/ ├── train/ │ ├── GOT-10k_Train_000001/ │ │ ├── 00000001.jpg │ │ ├── ... │ │ └── groundtruth.txt │ └── ... ├── val/ └── test/每个序列文件夹直接放图片和groundtruth.txt第一行是初始框后面每行对应每帧的目标框。TrackingNet不做训练集和测试集区分代码里直接读它的train文件夹需要在数据配置中指定对应路径。COCO需要按照标准格式组织因为代码里只用COCO的标注json文件来读取detection格式的目标框跟检测任务的preparation完全一致。3.2 数据路径配置四个数据集的路径需要在配置文件中显式指定。官方代码训练入口是train.py它会读取experiments/下的yaml配置文件。典型的配置片段像这样DATA: DATASET: mixed DATA_PATH: /your/path/to/datasets/ ...这里的DATA_PATH是数据集根目录的上一级也就是说代码会拼上lasot、got10k这些子目录名。我建议把所有数据集都放在同一个根目录下用软链接把不同盘符的数据统一到一处避免路径混乱。另外要注意代码里LaSOT用的是train子目录但有的LaSOT压缩包解压后目录结构里带有LaSOTBenchmark这样的外层文件夹需要把LaSOTBenchmark里的train和testing直接放到lasot/下否则路径对不上。4. 预训练模型与权重加载4.1 需要下载的文件清单Stark训练时需要的预训练权重主要有两类Backbone预训练权重默认ResNet-50官方在ImageNet上预训练但Stark的实际效果与骨干网络预训练方式关系很大。官方推荐使用经过MAE预训练或改良训练策略得到的ResNet-50权重效果比标准ImageNet权重有明显提升。你可以从官方仓库指定的链接下载文件名一般是类似resnet50_mae.pth或mae_pretrain_vit_base.pth这种。如果想直接复现论文指标更高效的是下载官方已经训练好的完整Stark模型权重文件大致是stark_s50_ep500.pth这种命名。下载后放在pretrained_models/目录下即可。需要注意官方权重文件的格式可能包含module.前缀这是在多卡训练时由DataParallel或DistributedDataParallel自动加上的名字空间。如果你直接加载到单卡模型会报UnExpected key的错误。解决方式很简单要么在代码里把state_dict的key去掉前缀要么用strictFalse加载后再手动转。4.2 在训练脚本里正确加载Stark的train.py在处理预训练权重时会先构建模型再调用load_pretrain函数更新模型参数。如果加载报错不建议直接改成strictFalse草草跳过因为有些key加载不上意味着backbone根本没有初始化成功训练出来的模型精度会直接崩。我写了一个通用的key对齐脚本可以参考import torch from collections import OrderedDict def load_weights(model, weight_path): checkpoint torch.load(weight_path, map_locationcpu) if state_dict in checkpoint: state_dict checkpoint[state_dict] else: state_dict checkpoint new_state_dict OrderedDict() for k, v in state_dict.items(): name k.replace(module., ) # 去掉多卡前缀 new_state_dict[name] v model.load_state_dict(new_state_dict, strictFalse) print(fLoaded weights from {weight_path})注意这里的strictFalse是为了忽略一些辅助预测头参数但backbone部分的key必须全部匹配上。可以在加载后打印一下missing keys和unexpected keys分别是什么确认missing里只有分类头或回归头相关参数时再放心往下走。5. 训练配置与实操记录5.1 训练入口与常用参数训练命令格式如下python train.py --config experiments/stark_s50/whatever_train.yaml完整的训练参数也会写在yaml里包括TRAIN: EPOCH: 500 BATCH_SIZE: 64 LR: 0.0001 BACKBONE_MULTIPLIER: 0.1 ...最关键的是BACKBONE_MULTIPLIER意思是backbone的学习率是主干网络的0.1倍这是为了在微调backbone时不破坏预训练特征。如果训练过程中发现loss下降很慢可以先检查这个值是否生效。训练过程中默认会使用混合精度但部分GPU型号对torch.cuda.amp支持不太好如果报混合精度相关的nan loss可以在配置里关掉AMP再试。5.2 损失函数与训练监控Stark的损失由两部分构成分类分支的损失和IoU预测分支的损失。分类分支负责判断搜索区域每个位置是否存在目标IoU分支直接预测预测框与对应真实框的IoU。用Tensorboard观察训练过程时重点看两条曲线loss_cls和loss_iou。正常情况下两者都会在训练初期快速下降大约50个epoch后趋于平稳。如果发现分类loss一直在0.7左右下不去大概率是训练数据里正负样本不均衡或者搜索区域中目标缺失太多导致的。5.3 单卡训练的显存优化方案如果你的卡显存不够优先改这几个参数按影响从大到小排列减小batch size比如从每卡16降到8或4。减小搜索区域尺寸把320改成288或256这个改动会直接影响最终跟踪精度但在复现阶段为了跑通可以接受。关闭部分数据增强里代价较高的操作比如随机仿射变换中的大旋转角度。如果代码版本支持gradient accumulation可以用累积步数模拟大batch。我实际在一张12GB显卡上训练过Stark的简化配置batch size设8搜索区域288backbone lr multiplier保持0.1其他不变。训练出来的模型在LaSOT上比官方完整配置低了约1.5个点但对于验证训练流程和调通代码来说完全够用。6. 测试流程与结果评估6.1 在测试集上运行模型测试入口在tracking/test.py。运行命令大概是python tracking/test.py --config experiments/stark_s50/test.yaml \ --dataset lasot \ --checkpoint pretrained_models/stark_s50_ep500.pth这里有一个容易踩坑的地方test.py读取参数的方式可能跟train.py不一致它更依赖命令行参数而非yaml。如果运行时报KeyError: SEQ_LEN之类的问题多半是没传全某些默认参数。测试过程会逐序列逐帧跑跟踪速度大概在几十帧每秒取决于GPU生成的结果默认存在一个results目录下每种数据集一个子文件夹里面每个序列对应一个txt文件。6.2 结果指标计算跟踪结果文件生成后下一步是计算成功率、精度等指标。官方评估脚本可能嵌套在tracking/eval.py里或者需要额外下载第三方评估工具。不同数据集的评估方式有差异LaSOT和TrackingNet用成功率AUC和精确度P两个指标。GOT-10k用平均重叠率AO和成功率SR。GOT-10k官方有严格的提交要求本地评估一般只能对照论文数值。计算AO时需要注意只有当预测框和真实框重叠率大于阈值时才认为是成功预测需要按官方脚本的阈值定义来走。评估步骤里最容易出问题的是结果文件名不对。比如LaSOT要求以序列名称命名有的脚本要求末尾加_post后缀跟代码里生成的文件名对不上时会报文件找不到。解决方案是在eval之前先打印前几个文件名人工核对一下格式。7. 常见问题快查与避坑心得7.1 问题速查表这里列出我试跑过程中遇到频率最高的4个问题和对应解法问题现象根本原因解决办法加载权重时报key不匹配多卡训练的module.前缀写脚本去掉前缀再严格加载训练loss突然变为NaN混合精度在某些GPU上不稳定关闭AMP或调整学习率测试时出现空txt结果测试脚本的读取路径不对检查数据集软链接是否正确启动训练时数据加载卡住数据集目录层级与代码预期不一致按第3节目录结构逐一核对7.2 几条实战建议第一建议先用官方权重跑通测试流程再动手训练。这一步能快速验证环境、数据、评估脚本全链路是否正常。如果官方权重测试指标都与论文差距过大优先排查数据集的groundtruth是否读取正确。第二训练前把数据集里每个序列都抽样一两帧人工可视化一遍。特别是COCO和LaSOT一起混训的时候如果某个数据集的bbox坐标格式理解错误模型会被带偏。第三修改模型结构后不要直接沿用原版的超参数。Stark里Transformer部分对学习率比较敏感如果改动encoder层数或维度建议把主干学习率降一个数量级观察几个epoch的loss曲线再决定是否继续。第四多卡训练时确保每个进程的seed一致否则数据增强的随机性会导致不同卡加载的样本分布不一致影响BN层的统计值更新。8. 经验结束语如果顺着这篇文的配置流程走一遍整个链路跑通其实压力不大。Stark这套代码放在2024年的视角看不算复杂但它在SOT领域的地位不可忽视——Transformer结构、时空建模、IoU预测分支这些设计后来被不少工作沿用和改进了。我自己的体会是复现这类经典算法最大的收获往往不是指标本身而是趁机把Transformer类跟踪方法的数据流、训练范式、评估链路完整摸了一遍。之后再去接触新出的MixFormer、OSTrack这类工作会发现很多模块都似曾相识。复现Stark的过程其实就是在为理解整个Transformer跟踪家族打地基。