PyTorch行人重识别test.py源码解析-开发者社区

PyTorch行人重识别test.py源码解析

在智能安防、跨摄像头追踪等实际场景中，如何从海量监控画面中准确“认出”同一个人，是行人重识别（Person Re-Identification, ReID）的核心任务。模型训练完成后，真正的考验才刚开始——我们不仅需要高效提取图像特征，还要确保这些特征在不同视角、光照和遮挡条件下依然具备强判别性。

本文深入剖析一个典型的test.py脚本实现，带你一步步拆解基于 PyTorch 的 ReID 推理流程。整个分析基于PyTorch-CUDA-v2.7 镜像环境展开，该环境预集成了 CUDA、cuDNN 和主流视觉库，开箱即用，极大简化了部署前的准备工作。

开发环境与运行方式

PyTorch-CUDA 基础镜像简介

当前使用的镜像是专为深度学习优化的容器化环境，版本为PyTorch 2.7 + CUDA 支持，适用于 NVIDIA V100、A100 及 RTX 30/40 系列显卡。它不仅仅是一个 Python 包集合，更是一套完整的工程化解决方案：

自动启用 GPU 加速，支持 DataParallel 和 DistributedDataParallel 多卡推理；
内置torchvision、torchaudio、tensorboard等常用扩展；
兼容 ONNX 导出与 TensorRT 部署链路；
提供 Jupyter Lab 和 SSH 两种交互模式，适配从实验调试到生产上线的不同需求。

开发者只需拉取镜像并启动容器，即可立即投入测试工作，无需再花费数小时配置依赖或解决版本冲突问题。

使用方式选择：Jupyter 还是 SSH？

如果你习惯可视化操作，推荐使用Jupyter Lab模式。通过浏览器访问界面后，可以上传数据集、编写代码块、逐行调试test.py中的关键函数。尤其适合初学者快速验证模型输出是否正常。

在执行任何 GPU 相关代码前，建议先设置设备可见性：

%env CUDA_VISIBLE_DEVICES=0

这能避免因多卡资源竞争导致的内存溢出问题。

而对于远程服务器或集群场景，则更适合采用SSH 登录方式进行开发。

ssh -p <port> username@<ip_address>

进入容器后，可使用vim或nano编辑脚本文件，并通过以下命令启动测试：

python test.py --batchsize 64

强烈建议配合tmux或screen工具运行，防止网络中断导致进程被终止。例如：
bash tmux new-session -d -s reid_test 'python test.py'

这种方式更适合长时间运行的大规模特征提取任务。

模型加载与数据准备

测试阶段的第一步，是从磁盘恢复训练好的模型权重，并构建标准输入流水线。

以常用的ft_net结构为例（基于 ResNet50 主干），其分类头在训练时用于身份分类，而在测试时则被移除，仅保留全局特征提取能力。假设我们在 Market-1501 数据集上训练完成，类别数为 751：

model_structure = ft_net(num_classes=751) model = load_network(model_structure)

其中load_network()函数负责加载.pth文件中的状态字典：

def load_network(network): save_path = os.path.join('./model', 'net_last.pth') network.load_state_dict(torch.load(save_path)) return network

注意：如果训练时使用了 DataParallel，保存的权重会带有module.前缀。此时需做兼容处理：

state_dict = torch.load(save_path) from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] if k.startswith('module.') else k new_state_dict[name] = v network.load_state_dict(new_state_dict)

否则会报错找不到对应层。

接下来是数据预处理部分。ReID 对输入尺寸有严格要求，通常统一缩放到(256, 128)，并进行标准化：

data_transforms = transforms.Compose([ transforms.Resize((256, 128), interpolation=3), # BICUBIC 插值 transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])

这里interpolation=3实际对应 PIL 中的Image.BICUBIC，相比双线性插值能保留更多纹理细节，对低分辨率监控图尤为重要。

然后分别构建 gallery 和 query 数据集：

image_datasets = { x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms) for x in ['gallery', 'query'] } dataloaders = { x: torch.utils.data.DataLoader( image_datasets[x], batch_size=opt.batchsize, shuffle=False, num_workers=4, pin_memory=True ) for x in ['gallery', 'query'] }

几个关键参数值得说明：

shuffle=False：保证图像顺序固定，便于后续结果对齐；
num_workers=4：利用多进程加速 I/O，但不宜设得过高以免引发内存瓶颈；
pin_memory=True：将张量锁定在主机内存中，使 GPU 可通过 DMA 快速读取，显著提升传输效率。

一切就绪后，即可开始无梯度推理：

with torch.no_grad(): gallery_features = extract_feature(model, dataloaders['gallery']) query_features = extract_feature(model, dataloaders['query'])

torch.no_grad()是推理阶段的标配上下文管理器，关闭自动求导机制后，内存占用可减少约 30%-50%，同时加快前向传播速度。

特征提取：精度与鲁棒性的双重保障

extract_feature()是整个test.py的核心函数，它的设计直接决定了最终检索性能。让我们来看其完整实现：

def extract_feature(model, dataloader): features = torch.FloatTensor().cuda() count = 0 for data in dataloader: img, _ = data n, c, h, w = img.size() count += n print(f"Processing {count} images...") ff = torch.FloatTensor(n, 512).zero_().cuda() for i in range(2): # 翻转增强 if i == 1: img = fliplr(img) input_img = img.cuda() for scale in ms: # 多尺度推理 if scale != 1: input_img = nn.functional.interpolate( input_img, scale_factor=scale, mode='bicubic', align_corners=False ) outputs = model(input_img) ff += outputs fnorm = torch.norm(ff, p=2, dim=1, keepdim=True) ff = ff.div(fnorm.expand_as(ff)) features = torch.cat((features, ff), dim=0) return features

这个函数融合了两项关键技术：测试时增强（TTA）和多尺度推理，共同提升了特征的一致性和泛化能力。

图像翻转增强（Horizontal Flip）

第一次循环使用原始图像，第二次将其水平翻转后再推理：

if i == 1: img = fliplr(img)

fliplr的实现如下：

def fliplr(img): inv_idx = torch.arange(img.size(3)-1, -1, -1).long().cuda() img_flip = img.index_select(3, inv_idx) return img_flip

这是一种轻量级的数据增强手段，在测试阶段也能生效。由于人体左右对称性较强，翻转后的姿态仍具有语义一致性。将两次推理结果相加，相当于做了简单的特征集成，能够缓解因局部遮挡或姿态变化引起的误判。

不过要注意，某些特殊结构（如带文字的背包）在翻转后可能失真，因此增强效果存在上限。实践中也有人尝试随机裁剪+翻转组合，但会增加计算成本。

多尺度推理（Multi-Scale Inference）

变量ms = [1.0, 1.1]表示使用两个尺度进行推理：

for scale in ms: if scale != 1: input_img = nn.functional.interpolate(...) outputs = model(input_img) ff += outputs

比如先以原图推理一次，再放大 10% 后重新推理，最后将两者特征累加。这种策略特别适用于监控摄像头分辨率不一的情况——小目标在放大后可能激活更强的响应。

虽然增加了约 1.5 倍的计算量，但在 mAP 指标上常能带来 1~2 个百分点的提升。对于追求极致性能的应用（如比赛提交），这是必不可少的一环。

当然，也可以进一步扩展为[0.9, 1.0, 1.1]更丰富的尺度集合，但边际收益递减明显，且显存压力增大。

特征归一化的重要性

fnorm = torch.norm(ff, p=2, dim=1, keepdim=True) ff = ff.div(fnorm.expand_as(ff))

L2 归一化使得所有特征向量落在单位超球面上，这样在后续计算余弦相似度时可以直接用内积代替，既高效又稳定。如果不做归一化，特征幅值差异可能导致距离度量失效。

此外，归一化还能缓解 Batch Normalization 层在不同 batch 下统计量波动带来的影响，提高跨批次一致性。

最终所有 batch 的特征通过torch.cat在样本维度拼接成完整矩阵，形成可用于检索的特征库。

图像元信息提取：评估的前提

要评估模型性能，仅有特征还不够，还必须知道每张图的身份标签（label）和摄像头编号（camera ID）。这些信息通常编码在文件名中，例如：

0001_c1_s1_0001.jpg → label=0001, camera=1 -1_c3_s2_0002.jpg → label=-1 (无效样本), camera=3

为此定义了解析函数get_id()：

def get_id(img_paths): labels = [] cams = [] for path, _ in img_paths: filename = os.path.basename(path) label = int(filename[0:4]) if filename[0:4] != '-1' else -1 cam = int(filename.split('c')[1][0]) labels.append(label) cams.append(cam) return cams, labels

调用方式简洁明了：

gallery_paths = image_datasets['gallery'].imgs query_paths = image_datasets['query'].imgs gallery_cam, gallery_label = get_id(gallery_paths) query_cam, query_label = get_id(query_paths)

这些元信息将与特征一同保存，用于后续计算 Rank-1、mAP 等指标。需要注意的是，同一身份出现在不同摄像头下的图像不应被视为正样本匹配，因此 camera ID 是去重和划分匹配空间的关键依据。

结果导出与后续分析

完成特征提取后，需将其保存为通用格式以便评估脚本读取。常用做法是导出为 MATLAB 兼容的.mat文件：

result = { 'gallery_f': gallery_features.cpu().numpy(), 'query_f': query_features.cpu().numpy(), 'gallery_label': gallery_label, 'query_label': query_label, 'gallery_cam': gallery_cam, 'query_cam': query_cam } scipy.io.savemat('pytorch_result.mat', result)

关键点在于：GPU 上的张量必须先通过.cpu()拷贝回主机内存，再转为 NumPy 数组才能被savemat序列化。否则会抛出类型错误。

生成的pytorch_result.mat可被evaluate_gpu.py或 MATLAB 脚本直接加载，用于：