AGW跨模态行人重识别实战:SYSU-MM01数据集80 Epoch训练复现指南
跨模态行人重识别(Cross-Modality Person Re-identification)是计算机视觉领域的前沿研究方向,旨在解决不同模态(如可见光与红外)图像间的行人匹配问题。AGW(Attention-Guided Wavelet)作为该领域的代表性方法,在SYSU-MM01数据集上实现了77.8%的Rank-1准确率。本文将提供从环境配置到模型训练的完整复现流程,帮助研究者快速验证论文结果。
1. 环境准备与数据集处理
1.1 基础环境配置
复现AGW模型需要以下环境依赖:
conda create -n agw python=3.8 conda activate agw pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html pip install numpy==1.20.3 scipy==1.7.1 opencv-python==4.5.3.56关键组件版本要求:
- CUDA 11.1
- cuDNN 8.0.5
- PyTorch 1.8.1
1.2 SYSU-MM01数据集下载与预处理
SYSU-MM01数据集包含491个行人的287,628张RGB图像和15,792张红外图像,采集自6个摄像头(4个RGB+2个IR)。数据集结构如下:
SYSU-MM01/ ├── cam1/ # RGB摄像头1 ├── cam2/ # RGB摄像头2 ├── cam3/ # IR摄像头1 ├── cam4/ # RGB摄像头3 ├── cam5/ # RGB摄像头4 └── cam6/ # IR摄像头2预处理步骤:
- 从 官方渠道 下载数据集
- 运行预处理脚本生成.npy文件:
python pre_processing_sysu.py --data_path /path/to/SYSU-MM01- 数据集将自动划分为训练集(395人)和测试集(96人)
注意:预处理时可能遇到路径错误,需检查data_manager.py中的路径设置
2. 模型代码与关键配置
2.1 代码库获取与结构
克隆官方代码库:
git clone https://github.com/mangye16/Cross-Modal-Re-ID-baseline cd Cross-Modal-Re-ID-baseline主要文件说明:
train.py: 主训练脚本test.py: 测试评估脚本model/agw.py: AGW模型实现data_loader.py: 数据加载器data_manager.py: 数据集管理
2.2 关键参数配置
在train.py中修改以下核心参数:
parser.add_argument('--dataset', default='sysu', help='dataset name: sysu or regdb') parser.add_argument('--lr', default=0.1, type=float, help='learning rate') parser.add_argument('--method', default='agw', help='method type: agw or baseline') parser.add_argument('--batch-size', default=8, type=int) parser.add_argument('--epochs', default=80, type=int) parser.add_argument('--gpu', default='0', type=str)推荐训练配置:
| 参数 | 值 | 说明 |
|---|---|---|
| batch_size | 8 | 根据GPU显存调整 |
| base_lr | 0.1 | 初始学习率 |
| optimizer | SGD | 动量0.9 |
| weight_decay | 5e-4 | L2正则化系数 |
| lr_schedule | [40,60] | 学习率衰减epoch |
3. 训练流程与技巧
3.1 启动训练
单卡训练命令:
python train.py --dataset sysu --lr 0.1 --method agw --gpu 0 --batch-size 8多卡训练建议:
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py \ --dataset sysu --lr 0.2 --method agw --batch-size 163.2 训练监控
训练过程中关注以下指标:
- 分类损失:应稳定下降至0.5以下
- 三元组损失:逐渐收敛至0.3左右
- Rank-1/mAP:验证集性能参考
典型训练曲线特征:
- 前20 epoch快速收敛
- 40 epoch后学习率衰减
- 60 epoch后微调特征
3.3 常见问题解决
显存不足:
- 减小batch_size(最低可设4)
- 使用梯度累积:
for i, (inputs, labels) in enumerate(dataloader): outputs = model(inputs) loss = criterion(outputs, labels) loss = loss / accumulation_steps loss.backward() if (i+1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()预处理报错:
- 检查
pre_processing_sysu.py中的路径 - 确保图像文件名格式正确
- 检查
4. 模型测试与性能验证
4.1 测试脚本运行
使用最佳模型进行测试:
python test.py --mode all --resume path/to/model.pth --gpu 0 --dataset sysu测试模式说明:
--mode all: 全搜索模式(All-search)--mode indoor: 室内搜索模式(Indoor-search)
4.2 性能指标解读
在SYSU-MM01上的预期结果:
| 模式 | Rank-1 | mAP | 备注 |
|---|---|---|---|
| All-search | 77.8% | 74.2% | 论文报告值 |
| Indoor-search | 82.3% | 79.1% | 论文报告值 |
实际复现可能出现的波动范围:
- Rank-1: ±1.5%
- mAP: ±2%
4.3 结果可视化
使用以下代码可视化检索结果:
import matplotlib.pyplot as plt def show_retrieval(query, gallery, indices): plt.figure(figsize=(15,5)) plt.subplot(1,6,1) plt.imshow(query) plt.title("Query") for i in range(5): plt.subplot(1,6,i+2) plt.imshow(gallery[indices[i]]) plt.title(f"Top-{i+1}") plt.show()5. 进阶优化策略
5.1 数据增强改进
推荐增强组合:
from torchvision import transforms train_transform = transforms.Compose([ transforms.Resize((256,128)), transforms.RandomHorizontalFlip(), transforms.Pad(10), transforms.RandomCrop((256,128)), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])5.2 模型微调技巧
非局部注意力调参:
nonlocal_params = { 'in_channels': 2048, 'sub_sample': True, 'bn_layer': True, 'use_self': True # 启用自注意力 }损失函数加权:
criterion = { 'id': CrossEntropyLabelSmooth(num_classes=395), 'tri': TripletLoss(margin=0.3), 'weight': [1.0, 1.0] # ID损失与三元组损失权重 }
5.3 跨模态对齐策略
AGW的核心创新点:
- 小波注意力引导的特征分解
- 模态共享特征空间投影
- 跨模态三元组挖掘
实现关键代码段:
class WaveletAttention(nn.Module): def __init__(self, in_channels): super().__init__() self.query = nn.Conv2d(in_channels, in_channels//8, 1) self.key = nn.Conv2d(in_channels, in_channels//8, 1) self.value = nn.Conv2d(in_channels, in_channels, 1) def forward(self, x): B, C, H, W = x.shape q = self.query(x).view(B, -1, H*W) k = self.key(x).view(B, -1, H*W) v = self.value(x).view(B, -1, H*W) attn = torch.softmax(torch.bmm(q.transpose(1,2), k), dim=-1) out = torch.bmm(v, attn.transpose(1,2)) return out.view(B, C, H, W)