余弦度量学习

简介

本仓库包含用于训练度量特征表示的代码，该表示可与 deep_sort 跟踪器 配合使用。相关方法在以下文献中有所介绍：

@inproceedings{Wojke2018deep,
  title={Deep Cosine Metric Learning for Person Re-identification},
  author={Wojke, Nicolai and Bewley, Alex},
  booktitle={2018 IEEE Winter Conference on Applications of Computer Vision (WACV)},
  year={2018},
  pages={748--756},
  organization={IEEE},
  doi={10.1109/WACV.2018.00087}
}

论文中使用的预训练模型可在这里找到。论文的预印本也可在此链接获取。此外，本仓库还提供了在 Market1501 和 MARS 数据集上训练模型的代码。

在 Market1501 上训练

以下说明假设您已将 Market1501 数据集下载至 ./Market-1501-v15.09.15 目录下。以下命令将使用上述论文中描述的余弦-softmax 分类器开始训练：

python train_market1501.py \
    --dataset_dir=./Market-1501-v15.09.15/ \
    --loss_mode=cosine-softmax \
    --log_dir=./output/market1501/ \
    --run_id=cosine-softmax

这将创建一个目录 ./output/market1501/cosine-softmax，用于存储 TensorFlow 检查点，并可通过 tensorboard 进行监控：

tensorboard --logdir ./output/market1501/cosine-softmax --port 6006

代码会从训练数据中分离出 10% 作为验证集。在训练的同时，运行以下命令以在验证集上计算 CMC 评估指标：

CUDA_VISIBLE_DEVICES="" python train_market1501.py \
    --mode=eval \
    --dataset_dir=./Market-1501-v15.09.15/ \
    --loss_mode=cosine-softmax \
    --log_dir=./output/market1501/ \
    --run_id=cosine-softmax \
    --eval_log_dir=./eval_output/market1501

该命令将无限期阻塞，以监控训练目录中保存的检查点，并对训练目录中的每个检查点在验证集上进行评估。评估结果将存储在 ./eval_output/market1501/cosine-softmax 中，可通过 tensorboard 进行监控：

tensorboard --logdir ./eval_output/market1501/cosine-softmax --port 6007

在 MARS 上训练

要训练 MARS 数据集，首先下载 评估软件，然后从 数据集官网 下载 bbox_train.zip 和 bbox_test.zip，并将其解压到评估软件目录中。以下说明假设它们分别存储在 ./MARS-evaluation-master/bbox_train 和 ./MARS-evaluation-master/bbox_test 目录下。训练可以使用以下命令启动：

python train_mars.py \
    --dataset_dir=./MARS-evaluation-master \
    --loss_mode=cosine-softmax \
    --log_dir=./output/mars/ \
    --run_id=cosine-softmax

同样，这将创建一个目录 ./output/mars/cosine-softmax，用于存储 TensorFlow 检查点，并可通过 tensorboard 进行监控：

tensorboard --logdir ./output/mars/cosine-softmax --port 7006

与 Market1501 类似，10% 的训练数据被划分为验证集。在训练的同时，运行以下命令以在验证集上计算 CMC 评估指标：

CUDA_VISIBLE_DEVICES="" python train_mars.py \
    --mode=eval \
    --dataset_dir=./MARS-evaluation-master/ \
    --loss_mode=cosine-softmax \
    --log_dir=./output/mars/ \
    --run_id=cosine-softmax \
    --eval_log_dir=./eval_output/mars

验证集上的评估指标可以通过 tensorboard 监控：

tensorboard --logdir ./eval_output/mars/cosine-softmax

测试

最终的模型测试是通过数据集作者提供的评估软件完成的。训练脚本可用于提取测试集的特征。以下命令会将 MARS 测试集的特征导出至 ./MARS-evaluation-master/feat_test.mat：

python train_mars.py \
    --mode=export \
    --dataset_dir=./MARS-evaluation-master \
    --loss_mode=cosine-softmax .\
    --restore_path=PATH_TO_CHECKPOINT

其中 PATH_TO_CHECKPOINT 是用于评估的检查点文件。需要注意的是，评估脚本需要进行少量调整以应用余弦相似度度量。具体来说，需将 utils/process_box_features.m 文件中的特征计算方式由最大池化（第 8 行）改为平均池化，并在文件末尾添加重新归一化的步骤。修改后的文件应如下所示：

function video_feat = process_box_feat(box_feat, video_info)

nVideo = size(video_info, 1);
video_feat = zeros(size(box_feat, 1), nVideo);
for n = 1:nVideo
    feature_set = box_feat(:, video_info(n, 1):video_info(n, 2));
%    video_feat(:, n) = max(feature_set, [], 2); % max pooling 
     video_feat(:, n) = mean(feature_set, 2); % avg pooling
end

%%% normalize train and test features
sum_val = sqrt(sum(video_feat.^2));
for n = 1:size(video_feat, 1)
    video_feat(n, :) = video_feat(n, :)./sum_val;
end

Market1501 脚本也包含类似的导出功能，其使用方法与 MARS 相同：

python train_market1501.py \
    --mode=export \
    --dataset_dir=./Market-1501-v15.09.15/
    --sdk_dir=./Market-1501_baseline-v16.01.14/
    --loss_mode=cosine-softmax \
    --restore_path=PATH_TO_CHECKPOINT

该命令将生成 ./Market-1501_baseline-v16.01.14/feat_query.mat 和 ./Market-1501_baseline-v16.01.14/feat_test.mat，供 Market1501 评估代码使用。

模型导出

要将您训练好的模型导出以便与 deep_sort 跟踪器 配合使用，请运行以下命令：

python train_mars.py --mode=freeze --restore_path=PATH_TO_CHECKPOINT

这将生成一个 mars.pb 文件，可直接提供给 Deep SORT 使用。同样地，Market1501 脚本也包含类似的功能。

cosine_metric_learning 快速上手指南

本工具用于训练度量特征表示，旨在与 deep_sort 追踪器配合使用。基于论文《Deep Cosine Metric Learning for Person Re-identification》，支持在 Market1501 和 MARS 数据集上训练模型。

环境准备

• 操作系统: Linux (推荐 Ubuntu 16.04+)
• Python: Python 3.x
• 深度学习框架: TensorFlow (需兼容原代码版本，通常为 TF 1.x)
• 硬件要求: 支持 CUDA 的 NVIDIA GPU (用于加速训练)
• 依赖库:
  • numpy
  • scipy
  • Pillow
  • matplotlib
  • tensorboard

建议通过 pip 安装基础依赖：

pip install numpy scipy pillow matplotlib tensorboard

安装步骤

1. 克隆仓库

   git clone https://github.com/nwojke/cosine_metric_learning.git
   cd cosine_metric_learning
   

2. 准备数据集

   • Market1501: 下载数据集并解压至 ./Market-1501-v15.09.15。
     • 下载地址：Market1501 官网
   • MARS: 下载 评估软件 及数据集 (bbox_train.zip, bbox_test.zip)。
     • 将解压后的 bbox 文件放入评估软件目录，结构应为：
       • ./MARS-evaluation-master/bbox_train
       • ./MARS-evaluation-master/bbox_test
     • 数据集地址：MARS 官网

3. (可选) 下载预训练模型 如需直接使用论文中的预训练权重，可从 Google Drive 下载。

基本使用

以下以 Market1501 数据集为例，展示最核心的训练与监控流程。

1. 开始训练

运行以下命令启动训练，采用论文描述的 cosine-softmax 分类器：

python train_market1501.py \
    --dataset_dir=./Market-1501-v15.09.15/ \
    --loss_mode=cosine-softmax \
    --log_dir=./output/market1501/ \
    --run_id=cosine-softmax

• 训练检查点将保存至 ./output/market1501/cosine-softmax。
• 代码会自动划分 10% 的训练数据作为验证集。

2. 实时监控训练过程

开启一个新终端，运行 TensorBoard 查看损失曲线和准确率：

tensorboard --logdir ./output/market1501/cosine-softmax --port 6006

然后在浏览器访问 http://localhost:6006。

3. 同步进行验证评估

在训练进行的同时，开启另一个终端运行评估脚本（自动监测新生成的检查点）：

CUDA_VISIBLE_DEVICES="" python train_market1501.py \
    --mode=eval \
    --dataset_dir=./Market-1501-v15.09.15/ \
    --loss_mode=cosine-softmax \
    --log_dir=./output/market1501/ \
    --run_id=cosine-softmax \
    --eval_log_dir=./eval_output/market1501

• 评估结果将保存至 ./eval_output/market1501/cosine-softmax。
• 可通过新端口监控评估指标：

  tensorboard --logdir ./eval_output/market1501/cosine-softmax --port 6007
  

4. 导出模型供 Deep SORT 使用

训练完成后，冻结模型以便在 deep_sort 追踪器中加载：

python train_market1501.py --mode=freeze --restore_path=PATH_TO_CHECKPOINT

• 请将 PATH_TO_CHECKPOINT 替换为具体的检查点文件路径。
• 该命令将生成 .pb 格式的冻结模型文件。

注: 若使用 MARS 数据集，请将上述命令中的 train_market1501.py 替换为 train_mars.py，并调整相应的 --dataset_dir 和 --log_dir 路径（参考原文 MARS 章节）。