行人重识别训练全解析：从表征学习到train_model.py实战

行人重识别系列写到第12篇，前面把数据集、评价指标、特征提取这些基础打完了，今天终于要进入最核心的部分：训练模型。这篇我重点讲表征学习这条路线下的训练脚本 train_model.py 该怎么写、每个模块为什么这么设计、以及我实际跑代码时踩过的那些坑。其实初看行人重识别训练，会觉得跟图像分类没什么两样：加载数据、前向传播、算损失、反向传播就完事了。但真正动手写过一遍的人都知道，ReID 的训练里藏着大量细节，比如身份均衡采样、BNNeck、三组损失的难样本挖掘、学习率预热这些，每一个都直接影响最终检索效果。这篇文章我按代码的组织顺序来讲，从数据准备到损失函数再到训练主循环，最后附上问题排查清单，争取让没写过 ReID 训练的读者看完就能照着复现一个可用模型。

1. 训练脚本到底在训练什么：先把表征学习的设计逻辑捋清楚

1.1 为什么第一份训练代码选“表征学习”

行人重识别要做的事情，简单说就是给定一张行人图，去另一堆跨摄像头的图里找到同一个人的其他图片。这里最关键的一步，是把每张行人图映射成一个“特征向量”，让同一个人的特征距离近、不同人的特征距离远。那怎么让网络学到这种特征呢？表征学习（也叫 IDE，ID-discriminative Embedding）提供了一条非常直接的思路：把每个行人身份当作一个类别标签，把训练过程当作一个图像分类任务来优化。

我选择先讲这个范式，是因为它是整个 ReID 训练的基石。后面你看到的各种改进，不管是换更强的骨干网络、加 CenterLoss 还是引入 CircleLoss，底层绝大多数还是围绕“分类损失让特征可分、度量损失让特征紧凑”这两个目标在做文章。先把表征学习的训练脚本吃透，后续接任何新模块都只是在这个骨架上加代码。而且从工程角度说，表征学习训出来的模型已经能应付很多实际检索场景了，我自己做的一些小项目最初版本就是拿这种模型顶上线的，效果并不差。

你可能会问，既然是分类任务，直接用现成的图像分类训练代码不行吗？答案是可以跑，但效果上不去。ReID 训练有两个天然特性是普通分类没有的：第一，训练集里的每个身份样本数量差异很大，有的身份拍了十几张，有的只有两三张，直接用随机采样，模型会被大样本身份带偏；第二，最终使用场景不是分类，而是检索，所以训练时不能只做“把身份分对”这一个任务，还要在特征层面加上距离约束。这两点直接决定了 train_model.py 里要引入身份均衡采样、三元组损失、BNNeck 这些组件。

1.2 train_model.py 的整体职责划分

我把一个完整的 train_model.py 拆成 6 个模块：配置管理、数据加载、模型构建、损失函数与优化器、训练主循环、验证与模型保存。这个拆分方式不是我凭空想出来的，是从多次实验里总结出来的。一开始我的训练脚本就是拿分类代码改的，所有参数硬编码在文件里，跑一个实验要手动改源码，日志输出也是想起来就 print 一下，最后实验结果乱成一团。后来把所有可调项收敛到 argparse 和配置文件里，又补上了结构化的日志和自动保存逻辑，整个训练流程才变得可控。

这 6 个模块的协作关系可以用一句话概括：配置模块提供所有超参数和路径，数据模块把原始图片变成张量，模型模块负责产出特征和分类预测，损失模块计算梯度信号，训练主循环把所有部件串起来执行前反向更新，验证模块定期跑评测并决定是否保存新权重。每个模块之间尽量保持低耦合，比如模型模块完全不感知数据怎么增强的，验证模块也不关心优化器用什么学习率。这样做的好处很实际：当你后面想换评估策略，或者想加一个新的损失项，只需要动对应的模块，其他部分不用重写。

2. 动手实现 train_model.py 前的关键准备：数据、代码与参数

2.1 数据集目录结构：第一步往往是翻车现场

先花点时间把数据准备讲清楚。ReID 领域用得最多的公开数据集是 Market1501，目录结构大概是这个样子的：

Market-1501-v15.09.15/ ├── bounding_box_train/ 训练集所有行人图 ├── bounding_box_test/ 测试集所有行人图（包含 query 和 gallery 的并集，但通常不用这个目录做评估） ├── query/ 待检索行人图 └── gt_query/ 手工标注的匹配关系（可选）

如果你用的是别的工具包，可能还会看到bounding_box_train/下每个身份一个子文件夹的组织方式，这种用 PyTorch 的ImageFolder直接读就行。Market1501 原版不是按子文件夹组织的，而是一个扁平目录，文件名本身携带标签信息，例如0002_c1s1_000451_00.jpg，下划线分割后，第一个字段是身份 ID，第二个字段是摄像头信息。实际写训练代码时，我更推荐自己写一个 Dataset 类去解析文件名，而不是直接用 ImageFolder，因为后面组织三元组和评估时，往往还需要同时知道摄像头 ID 和帧号。

另外有个容易踩的坑：训练集和测试集的身份是完全不重叠的。Market1501 里有 751 个身份用于训练，750 个身份用于测试，如果训练代码读取测试集图片参与训练，评估指标会虚高一大截，但实际落地时模型能力很差。我第一次跑实验就吃过这个亏，当时是用 ImageFolder 把所有图片一股脑读进来做训练，结果 rank1 高得离谱，后来才发现数据读错了。所以写数据加载代码前，先把训练集、query、gallery 三个集合各自单独建 Dataset 实例，并做一次数量校验，确认没有交集。

2.2 训练管线拆解：从原始图片到 loss 反馈

数据管线这一块，我习惯把它分成三个层次：原始文件到样本、样本到张量、张量到模型输入。

第一个层次是定义 Dataset 的__getitem__，它负责读取某一张图片的路径、身份标签和摄像头标签。第二个层次是在 Dataset 内部应用数据变换，这里要区分训练模式和测试模式。训练模式会用到随机水平翻转、随机裁剪、随机擦除、归一化，测试模式只做 resize 和归一化。第三个层次就是 PyTorch 的 DataLoader，负责把单个样本拼成 batch、多进程预取数据，配合后面要讲的 Identity Sampler。

训练阶段为什么要加这么多种数据增强？因为 ReID 数据集的规模普遍偏小，Market1501 训练集才一万多张图，而行人外观受遮挡、姿态变化、光照影响极大。靠随机水平翻转和随机擦除，可以让模型在有限的样本里见到更多样的情况，本质是正则化。这里特别说一下随机擦除：它在图上随机选一块矩形区域，把像素值换成一个随机数值（或者用 ImageNet 均值填掉），强迫模型不要过度依赖某个局部外观特征，比如衣服上一个醒目 Logo。别小看这个操作，在实际实验里，加上随机擦除往往能给 mAP 带来 2 到 3 个点的提升。

2.3 关键超参数选择逻辑：每个数字都不是拍脑袋定的

训练 ReID 模型有几个核心超参数，我直接给一组我常用且效果稳定的配置：

为什么 batch size 要设计成 16 个身份、每个身份 4 张图？这是为了同时满足两个目标。分类损失希望 batch 里有尽可能多样的身份，而三元组损失希望在同一个 batch 里出现同一个人的多张图，只有这样才能找到正样本对。如果只用随机采样，一个 batch 里某个身份可能只有一张图，根本组不成正样本对。所以采样器必须保证“同一个身份的多张图同时出现在一个 batch 里”，这个细节对训练效果的影响非常大。

优化器为什么用 SGD 不用 Adam，可能很多人不理解。Adam 收敛确实快，但在 ReID 这种小数据集上容易带上过拟合，最终泛化不如 SGD 加动量。SGD 需要更精细的学习率调节，所以配了 warmup 加余弦退火。warmup 的目的是让网络从预训练状态平滑过渡到目标学习率，前 10 轮的学习率如果直接拉高，容易破坏预训练权重；余弦退火则能让学习率在训练后期逐渐变小，减少参数在最优解附近的震荡。这一套组合下来，模型收敛质量和稳定性都有保障。

3. 核心代码实现：从数据加载到反向传播

3.1 数据增强与身份均衡采样

训练阶段的数据变换我一般这么写：

import torchvision.transforms as T train_transform = T.Compose([ T.Resize((256, 128)), T.RandomHorizontalFlip(p=0.5), T.Pad(10), T.RandomCrop((256, 128)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), T.RandomErasing(p=0.5, scale=(0.02, 0.4)) ]) test_transform = T.Compose([ T.Resize((256, 128)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

这个写法里有一个细节值得注意：先 Pad 再 RandomCrop，相当于做了一个随机平移裁剪，效果类似于训练时让行人物体在画面里稍微偏移一点，对位置变化的鲁棒性有好处。随机擦除放在 Normalize 之后，这跟某些开源代码的顺序不太一样，其实两种做法都能用，放到 Normalize 之后可以保证擦除区域的像素值分布跟正常数据一致，不会出现特别离谱的数值。第二个细节是测试集绝不加随机擦除和随机翻转，我已经见过好几个人在验证代码里带上了训练时的增强，导致测试指标惨不忍睹还找不到原因。

接下来是身份均衡采样器。这个不太适合用RandomSampler实现，我通常直接写一个继承torch.utils.data.Sampler的类：

import random from torch.utils.data import Sampler class RandomIdentitySampler(Sampler): def __init__(self, data_source, num_instances=4): self.num_instances = num_instances self.id_to_indices = {} for idx, pid in enumerate(data_source.pids): if pid not in self.id_to_indices: self.id_to_indices[pid] = [] self.id_to_indices[pid].append(idx) self.ids = list(self.id_to_indices.keys()) self.num_identities = len(self.ids) def __iter__(self): indices = [] shuffled_ids = self.ids[:] random.shuffle(shuffled_ids) for pid in shuffled_ids: pool = self.id_to_indices[pid][:] random.shuffle(pool) indices.extend(pool[:self.num_instances]) return iter(indices) def __len__(self): return self.num_identities * self.num_instances

注意这个写法是每次迭代把所有身份打乱，然后每个身份抽 4 张图，最终的迭代器长度是身份数乘以 4。DataLoader 拿到这个索引序列后，再按照 batch size 切分，实际上是先把 P 个身份（这里取决于洗牌后的顺序，默认是每 P 个身份一组）拼接成 64 张图组成一个 batch。我想说的是，__len__返回 16 × 4 = 64，只是一个近似值，所有身份总数如果恰好是 P 的整数倍时，每个 batch 正好是身份数 P = 16。实际工程里如果身份数不能被 16 整除，最后一个 batch 会混入不同数量的身份，这时候我建议干脆让 DataLoader 的drop_last=True，保证每个 batch 都满足三元组计算的结构要求。

3.2 模型构建：Backbone + BNNeck + 分类头

模型部分我采用的经典搭配是 ResNet50 做特征提取器，后面接一个 BNNeck，再接一层全连接作为分类头。如果直接拿分类模型来用，会发现特征分布非常乱，因为分类损失只需要把特征推到某个类别中心附近就行，并不保证同类特征紧密聚集。BNNeck 的由来就是解决“分类损失和三元组损失冲突”的问题。分类损失倾向于优化特征和分类器权重之间的余弦关系，而三元组损失要在欧氏空间里度量距离，两者对特征范数的要求不一样。BNNeck 的做法很简单：特征出来后接一个 BatchNorm 层，BN 之前的特征用来算三元组损失，BN 之后的特征再接全连接算分类损失。这样做等于把两个损失的优化空间分开，各自管各自的那部分特征变换。

代码上大致长这样：

import torch import torch.nn as nn from torchvision import models class ReIDModel(nn.Module): def __init__(self, num_classes, feat_dim=2048): super().__init__() base = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1) self.backbone = nn.Sequential(*list(base.children())[:-2]) self.global_pool = nn.AdaptiveAvgPool2d((1, 1)) self.bnneck = nn.BatchNorm1d(feat_dim) self.bnneck.bias.requires_grad_(False) self.classifier = nn.Linear(feat_dim, num_classes, bias=False) def forward(self, x): feat = self.backbone(x) feat = self.global_pool(feat) feat = feat.view(feat.size(0), -1) bn_feat = self.bnneck(feat) logits = self.classifier(bn_feat) return logits, feat, bn_feat

feat是 BNNeck 之前的原始特征，主要给三元组损失使用；bn_feat是经过 BN 后的特征，用来算分类 logits。这里有一个很细微但很重要的设置：bnneck.bias.requires_grad_(False)。因为 BN 层的作用就是把特征分布居中到零均值附近，如果 bias 还可学习，就失去了这个意义。还有 ResNet50 的最后一层全连接被我去掉了，最后的 stage 输出是空间特征图，经全局平均池化后得到 2048 维向量。你要是想改轻量骨干，比如 ResNet34，只需要换掉 backbone 部分，输出维度改成 512 就行，其他代码不用动。

3.3 两个损失的实现与融合

训练用的损失函数我一般同时启用标签平滑交叉熵和三元组损失。先看标签平滑交叉熵的实现：

class LabelSmoothCrossEntropy(nn.Module): def __init__(self, num_classes, epsilon=0.1): super().__init__() self.num_classes = num_classes self.epsilon = epsilon self.log_softmax = nn.LogSoftmax(dim=1) def forward(self, logits, targets): log_probs = self.log_softmax(logits) with torch.no_grad(): smooth_targets = torch.full_like(log_probs, self.epsilon / self.num_classes) smooth_targets.scatter_(1, targets.unsqueeze(1), 1 - self.epsilon + self.epsilon / self.num_classes) loss = (- smooth_targets * log_probs).sum(dim=1).mean() return loss

为什么要做标签平滑？因为 ReID 训练集每个身份只有少量样本，网络很容易对训练集上的分类任务过度自信，学到的特征在测试集上泛化很差。平滑之后，真实类别的概率不是 1，而是 0.9 左右，剩余概率摊到所有类别上，强迫网络不要死记硬背标签，本质是一种正则化手段。

三元组损失部分我采用 Batch Hard 策略。Batch Hard 的意思是，对 batch 里每个 anchor，找出与它同身份最远的一张图作为正样本，再找出与它不同身份最近的一张图作为负样本，然后计算三元组损失。这样做比随机采三元组更高效，因为总是挑最难的样本去优化。

class TripletLoss(nn.Module): def __init__(self, margin=0.3): super().__init__() self.margin = margin def forward(self, features, targets): dist = torch.cdist(features, features, p=2) bs = features.size(0) mask_pos = targets.unsqueeze(1).eq(targets.unsqueeze(0)) mask_neg = ~mask_pos # 排除自身 eye = torch.eye(bs, dtype=torch.bool, device=features.device) mask_pos = mask_pos & ~eye # 每个 anchor 最远正样本、最近负样本 hardest_positive = dist.masked_select(mask_pos).view(bs, -1).max(dim=1).values hardest_negative = dist.masked_fill(mask_neg, -1e8) hardest_negative = hardest_negative.masked_fill(~mask_neg, 1e8) hardest_negative = hardest_negative.min(dim=1).values loss = torch.relu(hardest_positive - hardest_negative + self.margin).mean() return loss

这段代码有个容易出错的地方就是掩码计算。masked_fill(mask_neg, -1e8)并不是把负样本统一置为 -1e8，而是为了让那些不需要参与计算的位置（比如同身份的位置）变成一个很大的负值，然后再把所有非负样本位置统一置为 1e8，最后取最小值。这么写虽然有点绕，但避免了循环，速度很快。实际使用中我建议先构造两个布尔掩码，分别表示正样本对和负样本对，再做填充，逻辑会更清晰。核心点在于，计算三元组之前，特征向量最好做一下 L2 归一化，不然距离尺度可能跟 margin 不匹配，导致损失项对梯度贡献不稳定。

两个损失怎么融合呢？我在训练主循环里直接相加，两项权重都是 1.0：

logits, feat, bn_feat = model(images) loss_ce = ce_loss(logits, pids) loss_tri = triplet_loss(feat, pids) loss = loss_ce + loss_tri

如果发现某项损失值异常大或异常小，可以分别乘以对应权重来调整。最开始的实验里我建议先保持 1:1，跑一版 baseline，再根据趋势做调整。因为两个损失的收敛速度不同，早期分类损失下降快，三元组损失下降慢，后期两者都比较稳定，这种自然互补的效果基本能达到预期。

3.4 主训练循环与验证评估

训练主循环的框架比较固定，我习惯这样组织：

best_mAP = 0.0 for epoch in range(start_epoch, epochs): adjust_learning_rate(optimizer, epoch) model.train() for batch_idx, (imgs, pids, _) in enumerate(train_loader): imgs, pids = imgs.cuda(), pids.cuda() logits, feat, bn_feat = model(imgs) loss = compute_loss(logits, feat, pids, bn_feat) optimizer.zero_grad() loss.backward() optimizer.step() if batch_idx % 20 == 0: logger.info(...) if (epoch + 1) % 5 == 0 or epoch + 1 == epochs: mAP, rank1 = evaluate(model, query_loader, gallery_loader) if mAP > best_mAP: best_mAP = mAP save_checkpoint(...)

学习率调整函数用余弦退火：

def adjust_learning_rate(optimizer, epoch, warmup=10, base_lr=3.5e-4, total_epochs=60): if epoch < warmup: lr = base_lr * (epoch + 1) / warmup else: progress = (epoch - warmup) / (total_epochs - warmup) lr = base_lr * 0.5 * (1 + math.cos(math.pi * progress)) for param_group in optimizer.param_groups: param_group['lr'] = lr

验证阶段要做的事是把 query 和 gallery 所有图片过一遍模型，取出特征向量，然后计算距离矩阵。这里我提醒一个常见误区：很多人验证时会把特征直接送进距离函数，忘了做归一化。由于训练阶段分类头使用 BNNeck，最终查询特征最好也先做 L2 归一化再算余弦距离。评估指标一般至少看两个：rank1 和 mAP。rank1 表示第一张检索图正确的比例，mAP 衡量整体排序质量。只盯 rank1 会忽略列表里后面出现了多少正确结果，所以这两个指标通常一起报。

4. 训练调试实录：我踩过的问题与排查方法

4.1 高频问题速查表

写训练代码的路上，没有谁能一次跑通。我把这些年遇到的高频问题、原因分析和解决方案整理成一个速查表，方便你排查。

这里我特别想展开说第一条。loss 不动，十有八九不是模型问题，而是数据问题。我自己曾经调了一个下午，发现 loss 在 0.1 附近徘徊，最后检查 Dataset 才发现，读文件时把 pids 全部传成了 0，模型只用一张图的一个类别做分类，自然训练不出判别性特征。调试的时候，最好先在第一个 batch 结束后打印 logits 的形状和 targets 列表，看是不是符合预期；再跑一个小迭代，把 loss 值打印出来跟正常范围对比。ReID 训练里交叉熵损失正常会在 3 到 5 左右起步，三元组损失在 0.5 到 1 左右，如果差距太大，基本可以确定哪里写错了。

4.2 我保存模型和记录实验的习惯

模型保存策略，我推荐同时保存四个东西：model_state_dict、optimizer_state_dict、epoch、best_mAP。只保存 model 权重的话，中断续训时优化器状态和当前 epoch 丢失，学习率又得重新算，很容易影响最终效果。另一个习惯是每个 epoch 或者每 5 个 epoch 做一次验证，并记录到本地日志文件。Market1501 的 query 只有 3000 多张图，gallery 不算大，全量验证一次很快，但如果你的数据集更大，可以把验证频率降到每 5 个 epoch 一次，节省训练时间。

我还有一个很多人忽略的习惯：训练前做一次数据完整性校验。具体来说，读取训练集所有文件路径，检查文件是否存在、图片能否正常解码、标签是否能正确解析。ReID 数据集经常有命名不规范的文件，比如有些图片损坏无法用 OpenCV 读入，如果训练到一半才遇到，往往白白跑了好几个 epoch。做一次预校验的代码量很小，大概几十行，但能避免那种“训了一天发现数据全读错”的灾难。

日志记录我把每一次实验的配置、超参数、评价指标、保存路径都写到一个 csv 或者 json 里。因为在调参的过程中，你会同时跑很多版实验，如果不记录，一周后根本分不清哪个配置对应哪份权重。我的做法是在训练脚本开头生成一个实验名，比如resnet50_market1501_sgd_cosine，所有输出文件都放到这个目录下，这样对比时只看目录名和日志就能快速定位。

最后再补一个个人经验：即使是最简单的表征学习训练，也不能完全放任不管，建议至少每 5 个 epoch 观察一次 loss 和验证指标的变化趋势。如果训练到第 40 轮，mAP 已经不涨了，rank1 还在缓慢上升，这种模型往往排序质量一般，我更倾向于参考 mAP，而不是 rank1，来决定是否继续训练。我的经验是，mAP 的波动比 rank1 更能反映模型泛化能力，实际部署中用户也不只关心第一张图对不对，更要看前几张检索结果的整体质量。

训练脚本这个东西，理论上讲起来很简单，代码加起来也就几百行，但真正让模型效果拉开差距的，往往就是数据划分、采样策略、损失组合、学习率规划这些不起眼的细节。这篇文章里给出的方案都是我自己反复跑过、验证过比较稳定的搭配，你当然可以在这个基础上继续改，比如把三元组损失换成 CircleLoss，或者在骨干网络后面加一个注意力模块。不过在那之前，先把当前的代码吃透，确保每一行训练日志的变化都解释得通，再去做很花哨的改进才更有意义。

如果这套 train_model.py 你已经能顺畅跑通并且复现出不错的指标，下一阶段就可以进入新的坑了，比如难样本挖掘策略怎么调、多模型融合怎么做、重排序怎么能把 mAP 再往上抬一抬。那都是后话，先把今天这尾鱼煎好再说。