★ FEATURED ARTICLE 今天想聊的是ICCV2021里那篇很有名的跟踪算法Stark全名是Learning Spatio-Temporal Transformer for Visual Tracking。我刚把环境跑通的时候光是版本组合就折腾了快两天。如果你正准备复现这个算法或者想在它的代码上做二次开发这篇应该能帮你把很多坑提前填掉。Stark做的事很简单就是在单目标跟踪里引入Transformer来建模模板和搜索区域之间的时空关系。传统方法靠在线更新模板或者多个分支堆叠Stark直接用一个时空Transformer把初始模板、动态模板和当前搜索区域放在一起做注意力编码最后用回归头输出目标框。无论你是想复现指标、学习跟踪网络结构还是做后续改进这套配置都是绕不开的起点。先说清楚一点我没有你本地的具体目录和显卡信息所以下面涉及的路径、命令和参数都是按我自己的配置经验整理出来的一版通用流程。你动手的时候只要把路径和版本号替换成自己环境里的真实值就行思路是一样的。1. 先搞清楚Stark到底在做什么1.1 标题拆解Spatio-Temporal Transformer 在跟踪里的含义Stark这个算法名字里的“Spatio-Temporal”是核心。拆开看“Spatial”指的是同一帧画面里目标模板和搜索区域之间的空间关系而“Temporal”则关注不同时刻帧之间的联系。传统跟踪器通常把模板固定成第一帧的目标最多做在线微调。Stark的做法不一样它在Transformer编码器里同时输入了三种信息初始模板第一帧给定的目标区域提供稳定的外观基准。动态模板跟踪过程中过去某几帧的预测结果用来应对遮挡、形变和光照变化。当前搜索区域当前帧中可能包含目标的大范围图像块。这样设计的好处是模型在推理时不需要频繁的在线更新也不需要复杂的模板记忆机制。注意力机制会自动决定当前帧应该更相信初始模板还是最近的动态模板。这个思路在当时很新也是后续很多Transformer跟踪方法的源头之一。1.2 Stark的基本工作流程如果你直接去看代码会发现Stark的pipeline其实比想象中干净。首先把模板区域和搜索区域分别裁出来缩放成固定尺寸例如模板128x128、搜索区域288x288。然后一起送进一个CNN骨干网络提取特征再把特征展开成序列丢给Transformer编码器做时空建模。最后分类头判断每个位置上有目标的概率回归头直接输出边界框的坐标偏移。整个流程里有两个容易看晕的点。第一个是“动态模板”的生成代码里面通常是用当前帧的预测框去裁剪上一帧或前几帧的目标区域再和初始模板一起编码。第二个是Transformer的输入是多个帧的特征拼在一起所以序列长度比普通分类任务长很多这也是显存占用高的主要原因之一。1.3 为什么这篇算法值得配置和复现Stark在几个主流单目标跟踪数据集上都取得了当时的领先结果尤其擅长处理长时跟踪中的目标再出现和大幅形变。这些场景在实际业务里特别常见比如视频监控里的行人跟随、无人机视角下的车辆跟踪、体育视频里的目标剪辑甚至智能录像中的关键目标锁定。更实际的价值在于它的代码结构是很多后续Transformer跟踪器的基础。你把这个仓库跑通了再去看后面的改进版本就能一眼认出哪些模块是继承Stark的哪些是后来加的。这也是我强烈建议你把配置过程走一遍而不是直接看论文的原因代码里的数据流比论文里的示意图具体得多。2. 配置环境前必须想清楚的问题2.1 我的配置环境和版本选型Stark属于典型的“论文代码”风格依赖版本卡得比较死。我实测下来比较稳妥的一套是操作系统Linux。Windows理论上也能跑但数据集路径、进程管理、共享内存这些地方会多出很多莫名其妙的问题。Python3.8。太新的版本比如3.10以上某些第三方库的预编译包可能跟不上。深度学习框架PyTorch 1.9左右配套的torchvision也要对应上。我有段时间用2.x跑结果一些接口改名导致代码直接崩。GPU驱动环境建议用支持更多并行计算能力的新版本驱动具体看你手里的显卡和官方依赖来选。其他依赖OpenCV、PyYAML、matplotlib、tqdm、pandas缺哪个装哪个。这套组合不是唯一答案但它的好处是跟论文原始代码的兼容性最好。如果你不想在环境上花太多时间直接用这个组合会省心很多。2.2 为什么版本组合这么敏感很多配置问题都出在Python和PyTorch的版本更新太快。原作者写代码时有些接口的用法在现在的版本里已经被改掉了。比如torchvision里的transform操作、torch.nn里某些初始化函数、还有分布式训练相关的API每个大版本都可能偷偷变化。Stark这种研究代码没有团队长期维护靠的就是老版本环境。所以我的建议是先别急着把环境升级到最新。跑通训练和测试之后你再决定要不要做迁移。一个稳定的老版本能帮你把代码层面的坑和版本层面的坑分开不然你是真的分不清报错是逻辑问题还是版本兼容问题。2.3 显存和硬件到底需要多高先说结论如果只是测试别人的预训练模型24G显存的显卡会很舒服甚至一些12G显存的卡也能跑只是batch size要调小。如果要从零训练单卡24G会非常吃紧大部分人是多卡分布式训练。Stark的显存消耗大头不在输入图像本身而在Transformer的注意力计算。搜索区域288x288的特征图展开成序列之后注意力矩阵的大小是序列长度的平方。分辨率每提高一点显存消耗会加速上涨。配置里如果默认输入尺寸是288而你想改成320可能直接就把显存顶爆了。这个后面我会详细说怎么调。3. 从零开始搭建Stark环境3.1 第一步创建干净的虚拟环境我非常不建议直接在系统Python里装依赖因为接下来你会反复安装不同版本的包早晚会冲突。建议用环境管理工具创建一个独立环境。这里以Linux下的命令行操作为例python -m venv starkenv source starkenv/bin/activate如果你更习惯用其他环境管理工具完全没问题。关键是新建一个独立的环境然后在这个环境里统一安装所有依赖。如果你是在服务器上操作记得确认好自己的Python版本是3.8左右不要直接用系统默认的Python 3.5或者3.12。3.2 第二步安装深度学习框架与基础依赖框架的安装命令会根据你的GPU计算平台版本不同而不同最靠谱的方式是去官方安装命令页面选择对应Python和计算平台版本复制粘贴。以PyTorch 1.9左右为例大概长这样pip install torch1.9.0 torchvision0.10.0装完之后建议先做一个小验证在Python里跑一下是否有GPU可用import torch print(torch.__version__) print(torch.cuda.is_available())如果输出False说明框架没有正确识别到GPU要么是版本装错要么是显卡的驱动和加速环境不匹配。这一步一定要先确认好不然后面所有训练脚本都会卡在设备选择上。之后再用pip安装其他依赖pip install numpy opencv-python pyyaml matplotlib tqdm pandasOpenCV版本太高有时候会有API改动如果后续报cv2.VideoCapture或者cv2.dnn之类的错误可以先降一个大的版本号试试。3.3 第三步下载源码和预训练权重Stark的代码是开源的。你直接在代码托管平台上搜索“STARK跟踪”或“Spatio-Temporal Transformer for Visual Tracking”找到官方仓库克隆到本地。这一步没什么难度但要注意仓库里是否带了子模块如果带了子模块需要额外执行一次子模块拉取命令否则某些目录是空的。预训练权重通常不会放进源码仓库而是放在网盘或开源模型库里。下载后请注意权重的文件名最好保持和官方一致。代码里经常用固定的文件名匹配权重乱改名会导致加载失败或指标异常。权重目录的位置尽量放在仓库根目录下的checkpoints或类似目录里。后面配置文件里会用到这个路径。下载完先校验文件大小。如果大小和官方标注差很远大概率是下载不完整重新下载就好。这一步最常见的坑是权重下载了但代码提示no checkpoint found。原因基本就是路径配置不对。你打开配置文件把权重绝对路径写进去或者按仓库要求把权重放到指定目录都能解决。3.4 第四步准备标准化数据集Stark的训练和测试用到的是单目标跟踪领域常用的几个数据集比如LaSOT、GOT-10k、TrackingNet以及COCO的子集。不同的数据子集对应不同的配置文件。你在测试阶段不一定需要全量下载可以先只准备一个测试集比如某个较小规模的测试集用来验证环境是否正常。数据集目录结构建议保持官方仓库的预期。一般来说每个数据集目录下会有train和test两个文件夹每个视频序列又包含一个img文件夹和若干个标注文件常见的有groundtruth.txt。如果你把数据集放在/data/tracking下配置里就写这个根路径后面代码会根据配置里的数据集名字去拼接子路径。如果你只是想先把测试跑通不要从零训练那么可以暂时不下载训练集。直接使用官方预训练权重在测试集上生成预测结果然后用官方评测脚本计算指标。这套流程晚点再补训练数据也来得及。3.5 第五步修改项目配置文件Stark的仓库里一般会有YAML格式的配置文件这是整个项目最核心的“面板”。它会告诉代码使用哪个数据集、模型参数怎么设、训练轮数、batch size、学习率等等。以我常用的一个配置为例核心逻辑大概是这样STARK: DATA: ROOT: /data/tracking DATASET: lasot MODEL: BACKBONE: resnet50 INPUT_SIZE: 288 TEMPLATE_SIZE: 128 USE_TEMPORAL: True TRAIN: EPOCH: 50 BATCH_SIZE: 16 LR_INIT: 0.0001 WORKERS: 4这里INPUT_SIZE是搜索区域的大小TEMPLATE_SIZE是模板区域的大小。USE_TEMPORAL控制是否启用动态模板分支这个开关非常值得玩。如果你只想看空间分支的效果可以把它的值改为False模型退化成只靠初始模板。很多人在做消融实验时就是改这个字段。修改配置时有三个建议路径一定要写绝对路径别用相对路径。因为项目里的脚本可能在不同目录下执行相对路径很容易指错。第一次跑测试时把WORKERS先改成0。这个参数是加载数据的线程数如果设得太大而系统资源不够会出现卡死现象。BATCH_SIZE不要一上来就按论文里的值设。先设成一个你的显卡能撑住的最小值跑成功后再慢慢往上加。4. 训练和测试的实操过程4.1 用预训练权重跑通测试流程拿到预训练权重后第一件事不是去训练而是先在测试集上验证环境是否正确。对于单目标跟踪测试流程一般是读取一个视频序列把第一帧的真实框作为初始模板之后逐帧预测目标框最后把所有预测结果保存下来。在Stark仓库里测试脚本一般会接收两个关键参数配置文件名和要跑的数据集名字。假如你想测试某个数据集的某个子集命令大概长这样python run_tracker.py --cfg stark.yaml --dataset lasot跑起来之后你会看到每个视频的跟踪结果一帧帧输出有时还会显示帧率和置信度。如果画面里目标框基本跟着物体走说明测试流程已经通了。之后再用仓库里的评估脚本计算成功率曲线和精确率。评估结果和论文报告的数字接近就说明你从权重加载、预处理到后处理全链路都没问题。这一步如果卡住了不要先怀疑算法先确认以下几点视频帧的读取顺序是否按文件名排序标注文件是否和视频一一对应搜索区域裁剪是否在原始图片上进行颜色通道是否是RGB而不是BGR。跟踪算法的预处理步骤非常琐碎一个通道顺序搞错了指标会掉一大截。4.2 从零训练的完整流程测试通完之后如果你想从头训练一个模型流程要比测试复杂得多。你需要下载训练集在配置里指定数据集根路径然后运行训练脚本例如python run_training.py --cfg stark.yaml训练和测试最大的区别在于数据采样。单目标跟踪的训练每次要采一个模板帧和一个搜索帧这两帧通常来自同一视频序列且相隔一定的帧数。这样模型才能学到时序上的形变和运动。所以训练数据的加载不是简单的一帧一张而是按序列成对地读取。如果代码里默认每个batch是“模板搜索区域”的配对那么显存占用也要按两份输入来计算。训练过程中的监控指标包括分类损失和回归损失。Stark的回归头一般用L1距离和GIoU做组合分类头用二元交叉熵。这些损失都会打印在日志里你可以根据趋势判断训练是否正常。正常情况是损失随着epoch增大而下降如果不降甚至上升优先检查学习率和batch size。学习率过大在初始阶段就会发散学习率过小又会让收敛变得非常缓慢。4.3 显存估算和batch size选择很多人第一次训练就直接按论文里的batch size填然后立刻遇到“CUDA out of memory”。其实论文里的batch size通常是在多张专业级显卡上用分布式训练得到的单卡复现时一定要调小。这里的调整不是靠感觉可以做一个简单估算。以搜索区域288x288、模板128x128为例输入图像本身每对大概是搜索区域 288 x 288 x 3 0.24 MB模板区域 128 x 128 x 3 0.05 MB加起来约0.3 MB但是真正吃显存的是中间特征图和注意力矩阵。Transformer编码器内部的注意力计算会把特征图的所有位置两两计算相关度。如果特征图分辨率是几十乘几十序列长度就到几千注意力矩阵的大小是序列长度的平方这才是主要消耗。我在一块24G显存的显卡上batch size从16降到8再从8降到4才勉强把训练流程跑通。建议你从batch size2开始如果显存占用不超过80%再逐步调大。你也可以在训练脚本里观察每步的显存峰值那个数值最直观。4.4 配置文件里几个容易被忽视的细节Stark的训练配置里有一个EPOCH用于控制训练轮次但如果你用的数据集很小开太多epoch不仅浪费时间还会过拟合。我的习惯是先用一半的轮次做一个快速验证看损失有没有下降趋势再决定是否跑完整个计划。还有一个值得注意的选项是是否使用动态模板。这个开关会直接影响模型的结构。开启USE_TEMPORAL后模型会额外加载一个动态模板相关的参数模块关闭后则不需要。如果你在预训练权重是开启状态的情况下强行用关闭状态的模型去加载通常会出现权重不匹配的报错。这时要么下载对应版本的权重要么修改加载逻辑来适配。另外如果配置文件里有PRETRAINED或PRETRAIN_PATH这类字段记得确保指向的是Stark论文中使用的骨干网络预训练权重而不是随便某个通用的图像分类权重。骨干的预训练权重影响很大用错会导致训练初期损失异常。5. 常见问题与排查技巧实录5.1 加载预训练权重时提示“键名不匹配”这个报错我遇到太多次了。直接原因是代码里用多卡训练保存的权重键名前面会带“module.”前缀而你的测试代码用的是单卡加载期望的键名没有这个前缀。解决方法有两种在加载权重时强制加一个逻辑把键名中的“module.”去掉。或者在测试脚本里也把模型包装成多卡形式。优先选第一种因为第二种等于为了加载权重而改变了模型结构容易引入额外的问题。加一行字符串替换就能解决的事情不要搞复杂。5.2 训练或测试时显存突然不够如果是在跑了几个step之后才报显存不足可能是因为前面的分支没有释放缓存或者是动态模板模块积累了多帧的中间特征。这种情况下你可以先调用缓存清理函数再重试。还不行的话就把batch size调小。如果显存一直维持在很高水平建议检查数据加载的线程数。线程数过大会让CPU预处理和GPU计算重叠得不好反而拉低训练速度。先把WORKERS设为0如果速度能接受再慢慢加到2或4。5.3 程序卡在数据加载阶段这个在服务器上很常见尤其是用了多线程数据加载之后。表现是训练日志停在“loading dataset”或“initializing dataloader”很久不动。排除路径问题之后最大的嫌疑是共享内存被占满。如果是在受限环境里运行需要调整共享内存的大小。如果是在本机运行可以把WORKERS降低或者直接设为0先确认是不是数据加载逻辑的问题。你也可以把数据集中的一两个序列复制到一个单独的小目录做一个mini测试集。最小化复现是排查这类问题最快的方式。5.4 指标和论文差一大截如果你的测试流程能跑通但算出来的成功率、精确率明显低于论文首先别急着怀疑权重损坏。最常见的原因是评估协议不一致。比如Stark在某个数据集上用的是“一次性评估”也就是整个视频只用第一帧初始化一次而你在测试时可能每帧都重新初始化或者每隔N帧重新初始化一次这会让指标虚高也可能因为实现问题反而偏低。另一个原因是颜色空间。OpenCV读出来的图像是BGR顺序PyTorch模型训练时用的是RGB。很多预处理代码会在读取后做通道交换但也有的代码没做。你可以自己输出一帧图像看看颜色是否正常如果是偏蓝偏红基本就是通道顺序问题。5.5 常见问题速查表现象可能原因处理方式权重加载失败键名带module前缀去掉module前缀再加载CUDA out of memorybatch size过大从batch size2开始调数据加载卡死共享内存不足调小WORKERS或调大共享内存图像颜色不对BGR和RGB混淆检查预处理通道转换指标偏低评估协议不一致确认初始化帧和评估模式训练损失波动学习率过高降低学习率或加warmup6. 我在实际配置中的几点体会这套配置折腾下来我最大的感受是Stark的算法结构其实不复杂真正的门槛在于环境依赖和数据格式。很多人一开始就想着跑训练结果在环境上反复碰壁。我的建议顺序是先测试后训练先用最小数据集跑通再谈复现论文指标。还有一个小技巧在源码目录里先跑一次自带的demo确认模型能加载、能推理。如果demo都跑不通后面所有操作都会叠加一层不确定性。先找到最短路径跑通再逐步往完整流程上扩展这个思路能用在你以后配置任何论文代码上。如果你后面想基于Stark做改进最值得动手的方向是换掉ResNet骨干网络或者把Transformer的层数减少用精度换速度。我试过把搜索区域从288降到256显存占用能下降一大截速度也快不少精度损失在可接受范围内。这种实验只需要改配置里的INPUT_SIZE改动成本很低非常适合作为你熟悉代码的入门练习。希望这篇配置记录能帮你顺利跑通Stark。最后再提醒一句遇到报错时把英文报错信息里的关键词复制到搜索引擎里查通常比对着代码盲猜有效得多。
阅读完成 · 觉得有帮助？
