如果你和我一样是在复现某篇行人重识别论文时第一次打开train_model(表征学习).py这个文件大概率会被里面一长串的 DataLoader、loss、schedule 和 evaluate 逻辑绕晕。我也是这么过来的。今天就把这个贯穿整个训练流程的核心脚本拆开来讲一遍重点聊每个模块为什么这么设计以及实际跑通一个表征学习模型时哪些地方最值得盯紧。作为本系列的第十二篇这篇文章聚焦的就是模型训练这最后一道工序——输入是配置、数据和模型输出是一份可用的权重文件和评价指标。1. 为什么ReID训练需要单独一个train_model.py1.1 表征学习范式在整个ReID技术里的位置行人重识别Person Re-Identification简称ReID要解决的事情说白了就是跨摄像头找人。同一个行人在摄像头A出现几分钟后在摄像头B又出现算法能不能判断这是同一个人。过去传统做法靠手工特征颜色直方图、纹理特征之类的天花板很低现在主流做法就是深度学习表征学习让卷积网络自己学一个能区分不同行人的特征空间。所谓表征学习核心思想是让模型把每一张行人图片映射成一个低维特征向量通常叫 embedding。训练的目标很直接同一个人的不同图片对应的特征向量应该很近不同人的图片特征向量应该很远。这个“近”和“远”靠什么衡量一般用欧氏距离或者余弦相似度。所以训练脚本的重点就是怎么设计损失函数把这种“同近异远”的关系压进网络参数里。这套思路听起来不复杂但真正实现起来有一堆工程细节数据怎么采样、batch怎么组织、损失怎么组合、学习率怎么调、验证阶段怎么算 rank-1 和 mAP。这些写在一个单独的训练脚本里就是为了把这些流程串起来让整个项目能一键跑起来。1.2 train_model.py 在项目里的角色和模块依赖在常见的ReID开源项目里代码组织一般是“配置 数据 模型 损失 评估 训练”这套结构。train_model.py通常处在最上层它自己不重复造轮子而是把下面这些模块拼装起来配置模块读取yaml或者py文件里的超参数比如epoch数、学习率、batch size、损失权重、数据路径数据模块负责读数据集、做数据增强、按PK策略采样输出一个又一个batch模型模块定义骨干网络通常是 ResNet50和后面的 embedding 层损失模块把分类损失、三元组损失组合起来返回总损失评估模块在验证集上计算 query 到 gallery 的距离得到 rank-1、rank-5、mAP 这些指标。train_model.py要把这些零件“装”成一台能跑的车同时还要负责打印日志、保存 checkpoint、控制随机种子这些都属于训练工程里的琐碎活。如果你只看论文不看代码很容易觉得训练就是loss.backward()加optimizer.step()实际上没那么简单。写这个脚本时我最深的体会是掉链子的地方几乎都在模块拼装的数据流上而不在模型结构本身。所以这篇博文我会从数据采样开始讲然后说模型和损失怎么搭配再讲训练循环和验证逻辑最后聊聊那些不写在论文里、但直接影响训练效果的坑。2. 训练脚本的数据链路从原始图片到有效batch2.1 数据集的目录结构和标签组织先看数据输入。以ReID领域最常用的 Market-1501 为例训练集和测试集的划分是固定的训练集有751个行人ID测试集有750个行人ID且训练和测试的ID完全不相交。这在ReID里是必须的——如果测试集里出现了训练时见过的行人那就成了“开卷考试”rank-1会虚高没有参考意义。图片文件名本身带标签比如0002_c1s1_000451_03.jpg这种格式前四位是行人ID第二位字段是摄像头ID。训练脚本要做的第一件事就是把文件名解析成“图片路径 行人ID 摄像头ID”的列表。数据集类通常长这样class Market1501Dataset(Dataset): def __init__(self, root, splittrain, transformNone): self.transform transform self.images [] self.pids [] self.camids [] for fname in os.listdir(os.path.join(root, split)): if not fname.endswith(.jpg): continue pid int(fname[:4]) camid int(fname[5]) self.images.append(os.path.join(root, split, fname)) self.pids.append(pid) self.camids.append(camid)这里有一个容易被新手忽略的点行人ID要重新映射成从0开始的连续编号。不要直接用0002这种原始编号去喂分类损失因为模型输出维度必须等于训练集ID数原始编号中间有空洞会导致维度对不上。我一般会先收集所有pid做一个pid2label的字典再转一遍。这个步骤虽然简单但漏掉会直接报错而且报错信息不容易定位。2.2 为什么必须用PK采样而不是普通随机采样普通图像分类的 DataLoader 用随机采样就够了因为每个batch里的图片彼此独立有没有同一个类别出现都不影响分类损失计算。但ReID的表征学习不一样三元组损失Triplet Loss要求一个batch里同时出现“锚点图、正样本图、负样本图”并且正样本必须是同一个人在不同摄像头下的图片。这里有个训练思路的区别值得展开说。交叉熵损失只看单张图片把行人ID分类对就行和batch里其它图片是什么没关系但三元组损失的训练信号完全来自batch内部的图片配对关系。如果用一个随机采样的batch很可能抽到16张图全是不同的人那三元组虽然也能强行构造但都是“简单负样本”网络很快就学不到东西了。更极端的情况是同一个人的图只出现了一张根本组不成三元组。所以ReID训练脚本里几乎都会自己写一个 PK采样器每个batch包含 P 个不同的人每人挑 K 张图片batch size P × K。常见设置是 P8、K4也就是一个batch里8个人、每人4张图。这样既能保证三元组是有效的又能让每个batch有足够多的负样本对比。我写过一个简化版的 PK sampler 放在自己的脚本里核心逻辑是先把每个ID的所有图片序号收集好然后按ID随机抽样class RandomIdentitySampler(Sampler): def __init__(self, data_source, num_instances4): self.data_source data_source self.num_instances num_instances self.num_pids_per_batch 8 self.index_dic {} for idx, (_, pid, _) in enumerate(data_source): if pid not in self.index_dic: self.index_dic[pid] [] self.index_dic[pid].append(idx) self.pids list(self.index_dic.keys()) def __iter__(self): indices [] for pid in self.pids: idxs self.index_dic[pid] if len(idxs) self.num_instances: idxs idxs * (self.num_instances // len(idxs) 1) random.shuffle(idxs) indices.extend(idxs[:self.num_instances]) return iter(indices)注意这里有个细节如果某个ID的图片数量少于K我会把它的列表复制几份再截取保证每个人都能贡献 K 张图。有些实现会选择直接跳过ID不足K的人但那样会让少数ID永远不进训练偏置比较大。我在实操中倾向用复制法因为ReID数据集长尾现象很常见ID间图片数差距能到好几倍。2.3 数据增强的细节训练和验证要区别对待数据增强这块训练阶段通常用随机翻转、随机裁剪、颜色扰动以及行人重识别论文里特别常用的 Random Erasing。Random Erasing 的作用是随机选一块矩形区域把像素抹掉迫使网络不要只依赖某个局部特征。比如行人可能被遮挡、背包挡住身体这个增强能模拟这种遮挡情况提升模型泛化能力。关键坑在于验证和测试阶段不要做随机增强。验证集只需要做统一缩放到固定尺寸比如 256×128和归一化任何随机性都会让特征不一致导致同一张图每次推理出来的 embedding 都不稳定rank 指标忽高忽低。这个道理其实和分类任务一样但ReID场景下更敏感因为后面算距离矩阵时所有特征要可比增强差异直接污染距离。3. 模型与损失设计embedding 是怎么学出来的3.1 骨干网络的选择与 embedding 层设计现在ReID领域提到表征学习主干基本都默认ResNet50预训练模型。原因很简单ImageNet预训练权重能提供不错的底层视觉特征ReID数据集本身只有几万张图靠这些数据从头训一个深层网络很容易过拟合不如微调预训练权重划算。网络结构上通常把 ResNet50 最后一个全连接分类层去掉保留到全局平均池化输出2048维特征。然后在这个特征后面再接一个全连接层把维度降到256这个256维向量就是最终用于算距离的 embedding。为什么降到256而不是直接用2048一方面低维特征计算距离更快、存储更省另一方面论文和实验都表明高维特征并不一定带来更好的检索效果去冗余之后反而更鲁棒。这里我要特别提一下 BNNeck 这个设计。它是某篇ReID论文提出的小技巧在被大量开源实现采用后成了标配。做法很简单从骨干出来的特征经过一个全连接层得到256维embedding之后再接一个BatchNorm层。这个BN层的输出用于计算交叉熵分类损失而BN之前的原始特征用于计算三元组损失。为什么要分开因为两个损失对特征分布的要求是矛盾的。交叉熵希望特征分布有较大的类间角度差异三元组损失希望类内特征紧凑直接让同一个特征同时满足两个目标会互相拖累。BN层能把特征分布拉到一个比较集中的区域分类损失用起来更稳定而三元组损失用BN前的特征保留了更好的相对距离信息。3.2 交叉熵和三元组联合损失为什么能互补损失设计是整个训练脚本里最核心的一环。目前主流做法是交叉熵损失 三元组损失联合总损失等于两者加权相加。交叉熵损失在ReID里就是把它当作一个多分类问题把每个行人ID当成一个类别网络输出每个ID的概率然后算交叉熵。这个损失的好处是收敛快、稳定早期训练能迅速把网络拉到“能区分不同人”的状态。三元组损失的目标更细它直接优化特征空间中的相对距离。给定一个锚点 anchor一个正样本 positive同一个人一个负样本 negative不同的人希望 anchor 到 positive 的距离 一个间距 margin 仍然小于 anchor 到 negative 的距离。用公式表达就是d(anchor, positive) margin d(anchor, negative)margin通常取0.3。如果这个条件不满足就产生梯度把正样本拉近、负样本推远如果已经满足了这个三元组的损失就是0不参与更新。这会导致一个问题训练后期简单三元组大量出现损失几乎不变网络提升变慢。所以后来的实现普遍用难样本挖掘——在一个batch内选出距离最远的正样本和距离最近的负样本构造三元组逼网络处理最难的case。我用的方案是batch内所有正负样本对都算一遍然后挑最难的这种“batch-hard”策略代码简单而且效果好。交叉熵和三元组配合使用的好处我用一个比喻来解释交叉熵像老师给你划重点告诉你这些人是哪一类三元组像体育教练训练肌肉记忆让你形成“和自家人站一起、和陌生人保持距离”的本能。少了前者训练前期会非常慢少了后者学出来的特征虽然能分类但检索时排序质量不行——因为检索靠的是距离排序不是分类准确率。损失组合代码一般长这样criterion_ce nn.CrossEntropyLoss() criterion_triplet TripletLoss(margin0.3) logits model(images) # 分类输出 embedding model.feature # BN之前的256维特征 loss ce_weight * criterion_ce(logits, pids) triplet_weight * criterion_triplet(embedding, pids)权重设置上我常用的比例是 ce_weight1.0、triplet_weight1.0两个损失各一半。有些论文会把三元组权重调低一点比如0.5但实测下来在Market-1501上差别不大真正影响大的是学习率和数据采样。3.3 优化器与学习率调度warmup cosine 是标配ReID训练里最经典的学习率方案是前10个epoch做warmup把学习率从很小慢慢升到初始值之后用余弦退火cosine annealing让学习率平滑降到接近0。为什么要warmup因为模型刚加载预训练权重时后面的BN层和新增的embedding层都是随机初始化的如果一开始就用大学习率容易把这些随机层一下冲坏损失剧烈震荡。先小步走一会儿让这些层适应数据分布再进入正常训练节奏。优化器我一般用Adam初始学习率3.5e-4weight decay 5e-4。这个配置在多个ReID开源项目里被反复验证过属于“能用且不容易崩”的起点。如果你用的是纯SGD学习率要相应调大一些比如0.01但SGD对超参数更敏感调起来更费劲。对新手来说我建议直接抄Adam这套配置等跑通了再根据自己的数据集去微调。4. 训练主循环与验证逻辑哪些步骤决定最终指标4.1 一个epoch里到底发生了什么训练主循环本身并不神秘就是标准的“数据加载—前向—算损失—反向—更新参数”流程。但有几个细节值得展开。第一个是计算设备。ReID输入分辨率一般是256×128batch size 64 时单张GPU显存占用大约10GB左右所以常见做法是用单卡训练。多卡也不是不行但bn的同步和损失聚合需要额外处理我后面会专门讲。第二个是梯度累积。如果你显存只够跑 batch size 16但想模拟 batch size 64 的效果可以用梯度累积每4个batch才做一次optimizer.step()中间只累加梯度。这个技巧在调batch size受限时很好用但要小心学习率的换算关系batch size翻倍后学习率也可以适当上调不然收敛速度会变慢。第三个是统计丢失。ReID训练脚本里我习惯每个epoch打印四行信息损失值、当前学习率、batch耗时、累计epoch耗时。损失值是最直观的信号如果发现总损失在前几个epoch不降反升先检查是不是数据标签解析错了而不是急着调超参。我排查过好几次最后发现是PID重新映射时字典顺序打乱了导致损失张量里的标签和真实身份对不上。一个最小可用的训练循环大概长这样for epoch in range(start_epoch, max_epoch): model.train() for batch_idx, (imgs, pids, camids) in enumerate(train_loader): imgs, pids imgs.cuda(), pids.cuda() logits, embedding model(imgs) loss ce_loss(logits, pids) triplet_loss(embedding, pids) optimizer.zero_grad() loss.backward() optimizer.step() # 打印日志、更新学习率...这段代码虽然看着简单但实际项目里还有不少周边逻辑每个epoch结束后要保存当前权重、要跑一次验证、要记录最好的rank-1对应的权重还要支持从断点恢复训练。这些看起来是“工程细节”但缺了任何一个训练过程都会变得很难受。4.2 验证阶段rank-1 和 mAP 是怎么算出来的ReID的验证和分类任务完全不一样。分类任务验证只要在验证集上算一个准确率就完了ReID要模拟真正的检索场景。具体做法是把所有查询图片的embedding算出来构成query特征矩阵把所有候选库图片的embedding算出来构成gallery特征矩阵然后计算query和gallery两两之间的余弦距离。对每个query把所有gallery图片按距离从小到大排序看真正的同一个人排在第几位。排第一就是排名指标里的 rank-1排在前5就是rank-5。mAP则更严格它不光看排没排到第一还要看所有正确匹配在整个排序列表中的位置分布相当于对排序质量做整体评估。这里有个Market-1501特有的坑gallery里存在“同一行人在同一个摄像头下的图片”理论上这种样本不算真正的匹配。为什么因为查询图片本身也来自某个摄像头如果gallery里出现了同一个摄像头下的同一个人图片基本就是同一张图的重复帧计算距离太容易匹配上会让指标虚高。所以验证代码里要过滤掉这些样本规则是gallery中如果出现和当前query同ID且同摄像头ID的图片直接跳过不参与排序。我在代码里实现排序距离矩阵的简化版如下def compute_distance_matrix(query_feat, gallery_feat, metriceuclidean): # 两个矩阵都是 (N, 256) 维 m, n query_feat.size(0), gallery_feat.size(0) dist torch.zeros(m, n) for i in range(m): dist[i] torch.pow(query_feat[i] - gallery_feat, 2).sum(dim1) return dist实际项目里通常会一次算完所有距离用矩阵广播避免循环q2, g2 (query_feat ** 2).sum(dim1, keepdimTrue), (gallery_feat ** 2).sum(dim1, keepdimTrue) dist q2 g2.T - 2 * torch.mm(query_feat, gallery_feat.T)然后按行排序得到索引再结合query和gallery的标签列表统计每个query的匹配位置最终算出rank-1、rank-5和mAP。4.3 多卡训练与混合精度能用但不是必须的关于训练加速有两件事值得说。第一是混合精度训练。ReID模型本身不算大但训练时如果需要跑很多epoch混合精度还是能省不少时间。PyTorch自带的torch.cuda.amp已经非常成熟只需几行代码就能开启。我实测在单卡上能带来大约20%~30%的加速显存占用也能降四分之一左右。但要注意混合精度下如果出现 loss 变成 NaN先别慌多半是某个算子在fp16下数值溢出可以把这个算子排除在amp之外或者加大 loss scale 的初始值。ReID这种纯卷积网络还好基本不需要额外处理。第二是多卡DataParallel。ReID不在batch里做同步BN的话多卡训练本质就是把不同子batch分到不同卡上最后汇总梯度和单卡在数学上等价。但BN如果分别统计每张卡上的batch都只有原来1/N的大小均值方差估计不准。ReID里特征质量受BN影响比较大所以我如果要用多卡会优先用同步BN方案把所有卡上的统计量聚合。不过说实话对大多数复现场景而言数据集也没大到必须多卡训练的程度单卡跑一天左右就能出结果。与其花时间折腾多卡不如把时间花在调参上。5. 训练里最容易出现的坑和我的调参心得5.1 BN 层的坑训练和测试行为不一致BN在ReID训练里的重要性远高于分类任务。因为BN的统计量是随batch变化的训练时一个batch里有P个不同IDBN能比较准确地估计当前分布但推理时如果一次只输入一张图BN的统计量会和训练时不一致导致特征分布漂移。具体到实践有两个注意点。一是如果你想在微调场景下冻结主干参数、只训练embedding层那么冻结BN是个必要操作。因为主干虽然被冻结了但BN统计量还是会随每个batch的数据更新这会导致前向输出不稳定。ReID里经常有人抱怨“训练loss正常但验证指标很差”很可能就是冻结不完全导致BN统计量震荡。二是保存checkpoint时要同时保存模型参数、优化器状态和epoch数。如果只存模型参数中途断点恢复后学习率调度会从头开始warmup再跑一遍明显影响最终精度。5.2 epoch 数和验证频率怎么权衡ReID训练不像分类那样epoch越多越好。因为三元组损失本质上是在一个batch内部构造困难样本如果训练过度模型会对训练集的特定困难样本过拟合验证集上的rank-1不升反降。我见过最典型的走势是前40个epoch mAP稳步上升50到70个epoch继续缓涨过了80个epoch之后开始波动甚至下降。所以我的建议是每个epoch结束都跑一次验证直接把验证集上的rank-1作为模型选择的依据——哪个epoch的验证指标最高就保存哪个checkpoint而不是傻傻地训完固定epoch就取最后一个。代码里我用一个变量记录最佳mAP每次验证后判断是否更新更新就覆盖保存best模型。5.3 随机性控制和参考配置复现ReID论文时随机种子对结果影响很大。因为PK采样和Triplet的难样本挖掘都依赖随机数不同随机种子下同一个代码配同一个数据集最终rank-1可能相差1到2个百分点。这个波动在论文报告里通常被平均掉但你自己复现时会觉得很奇怪。所以训练脚本开头一定要统一设置import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)此外torch.backends.cudnn.deterministic True会让卷积计算变得更确定但会牺牲一些速度如果只是训练不追求逐位复现可以不开只把随机种子固定住就够了。最后给一份我实际跑过的参考配置直接照着设基本不会出大问题参数取值说明骨干网络ResNet50 ImageNet预训练主流选择输入尺寸256×128按行人长宽比训练epoch60~12060可先验证pipelinebatch size64P16, K4 或 P8, K8显存不够用梯度累积优化器Adam初始lr 3.5e-4weight decay5e-4与lr配套warmup epoch10前10个epoch线性上升学习率调度cosine退火到0也可最后固定小lr损失权重CE1.0, Triplet1.0可微调但不必大改三元组margin0.3调大偏向难样本验证频率每个epoch一次记录best模型5.4 几个我从踩坑里总结的小习惯先说日志。训练脚本的输出信息一定要包含当前epoch、损失、学习率、数据加载耗时、验证rank-1和mAP。这些信息写在一行里方便扫一眼就知道训练是否正常。如果损失下降正常但验证指标不动问题往往在数据增强或者损失权重上而不是在优化器。再说checkpoint管理。我会固定保存两个文件一个是当前epoch的latest模型覆盖式保存一个是最佳验证结果对应的best模型。这样即使程序跑到一半崩了也能从latest恢复继续训练而且不会丢掉之前的最佳结果避免后期过拟合导致best模型被误覆盖。最后说一个容易被忽略的点如果复现某篇论文发现指标怎么都差一截先不要怀疑模型结构先检查数据加载是否正确特别是PID编号有没有对齐。我踩过的坑里至少有一半是这种组件级的低级错误造成的。我在实际使用中发现把train_model.py从“能跑”打磨到“跑得明白”真正花时间的不是那一百多行训练循环而是理解数据采样、损失组合和验证协议之间的关系。建议你拿到任何一套ReID代码先别急着跑先按这个顺序把脚本读一遍数据集类 → PK采样器 → 模型返回哪些输出 → 损失怎么算 → 验证时怎么过滤同镜头样本。这条链路通了后面改模型、换数据集、调超参就都有了底气。
