通过同时进行聚类和表示学习实现自标签化

🆗🆗🎉 新增模型（2020年8月20日）：添加了标准的 SeLa 预训练 torchvision ResNet 模型，使加载更加方便；增加了使用更优 MoCov2 数据增强的基线模型（无监督学习性能约 69%）；为 ImageNet 的“无监督聚类”任务添加了 K=1000 的评估。

🆕✅🎉 更新代码：2020年4月23日：修复了若干 bug + 添加了 CIFAR 数据集的代码 + 为 resnet 和 alexnet 提供了评估功能。

请查看我们的 博客文章，其中提供了简短的非技术性概述以及我们聚类结果的交互式可视化。

自标签化

本代码是 ICLR 2020 论文《通过同时进行聚类和表示学习实现自标签化》的官方实现，论文链接：https://openreview.net/forum?id=Hyx-jyBFPr。

摘要

将聚类与表示学习相结合，是深度神经网络无监督学习中最有前景的方法之一。然而，若以简单粗暴的方式直接结合，往往会导致病态的学习问题，并产生退化解。在本文中，我们提出了一种新颖且具有理论基础的学习框架，以解决这些问题。该方法的核心思想是最大化标签与输入数据索引之间的互信息。我们证明，这一准则可以将标准的交叉熵最小化扩展为一个最优传输问题，而我们则利用 Sinkhorn-Knopp 算法的高效变体，能够快速求解包含数百万张图像和数千个类别的大规模问题。最终得到的算法能够自动为视觉数据打上标签，从而在无需人工标注的情况下训练出极具竞争力的图像表示模型。我们的方法在 SVHN、CIFAR-10、CIFAR-100 和 ImageNet 数据集上，分别针对 AlexNet 和 ResNet-50 模型，均取得了当前最先进的表示学习性能。

结果概览

	NMI(%)	aNMI(%)	ARI(%)	LP 准确率 (%)
AlexNet 1k	50.5	12.2	2.7	42.1
AlexNet 10k	66.4	4.7	4.7	43.8
R50 10x3k	54.2	34.4	7.2	61.5

使用更优数据增强（均为单裁剪）

	标签准确率	NMI(%)	aNMI(%)	ARI(%)	LP 准确率 (%)	模型权重
Aug++ R18 1k (新)	26.9	62.7	36.4	12.5	53.3	此处
Aug++ R50 1k (新)	30.5	65.7	42.0	16.2	63.5	此处
Aug++ R50 10x3k (新)	38.1	75.7	52.8	27.6	68.8	此处
(MoCo-v2 + k-means**, K=3k)		71.4	39.6	15.8	71.1

• “Aug++” 指的是 SimCLR 中使用的更优数据增强策略，源自 MoCo-v2 仓库，但我仍然只训练了 280 个 epoch，并按照 CMC 的方式进行了三次学习率衰减。
• 如果采用 MLP 或训练 800 个 epoch（我只训练了 280 个），并参考 SimCLR、MoCov2 和 SwAV 的做法，仍有进一步提升的空间。
• **MoCo-v2 使用 800 个 epoch、MLP 和余弦退火学习率调度。在 MoCo-v2 中，我会对平均池化的特征（经过 MLP 头后性能基本一致）运行 k-means 聚类（K=3000），以获得 NMI、aNMI 和 ARI 指标。
• 上述模型均使用标准的 torchvision ResNet 主干网络，因此加载变得非常简便：

import torch, torchvision
model = torchvision.models.resnet50(pretrained=False, num_classes=3000)
ckpt = torch.load('resnet50-10x3k_pp.pth')
model.load_state_dict(ckpt['state_dict'])
pseudolabels = ckpt['labels']

• 关于改进潜力的说明：仅使用“Aug++”数据增强时，我在训练 200 个 epoch 后获得了 67.2% 的 LP 准确率。而同样使用“Aug++”但训练 200 个 epoch 的 MoCo-v2，其 LP 准确率仅为 63.4%。

我们方法发现的聚类

按纯度排序

随机排列

图中边的颜色编码了真实的 ImageNet 类别标签（这些标签并未用于训练）。您可以在 这里 查看所有聚类。

系统要求

• Python >3.6
• PyTorch > 1.0
• CUDA
• Numpy、SciPy
• 另外，请参阅 requirements.txt 文件
• （可选）TensorboardX

运行我们的代码

运行 AlexNet 的自监督训练，使用以下命令：

$./scripts/alexnet.sh

或者训练 ResNet-50，使用：

$./scripts/resnet.sh

注意：你需要指定数据集目录（它期望的格式与 ImageNet 相同，包含 “train” 和 “val” 文件夹）。此外，你还需要为代码分配足够的 GPU，以便在 GPU 上存储激活值。否则，你将需要使用 CPU 版本，但其速度会显著慢得多。

无监督训练代码 main.py 的完整文档如下：

用法: main.py [-h] [--epochs EPOCHS] [--batch-size BATCH_SIZE] [--lr LR]
               [--lrdrop LRDROP] [--wd WD] [--dtype {f64,f32}] [--nopts NOPTS]
               [--augs AUGS] [--paugs PAUGS] [--lamb LAMB] [--cpu]
               [--arch ARCH] [--archspec {big,small}] [--ncl NCL] [--hc HC]
               [--device DEVICE] [--modeldevice MODELDEVICE] [--exp EXP]
               [--workers WORKERS] [--imagenet-path IMAGENET_PATH]
               [--comment COMMENT] [--log-intv LOG_INTV] [--log-iter LOG_ITER]

PyTorch 实现的自标签方法

可选参数:
  -h, --help            显示此帮助信息并退出
  --epochs EPOCHS       训练轮数
  --batch-size BATCH_SIZE
                        批量大小（默认：256）
  --lr LR               初始学习率（默认：0.05）
  --lrdrop LRDROP       每隔一定轮数将学习率乘以 0.1（默认：每 150 轮）
  --wd WD               权重衰减指数（默认：-5）
  --dtype {f64,f32}     SK 算法的数据类型（默认：f64）
  --nopts NOPTS         伪选项的数量（默认：100）
  --augs AUGS           数据增强级别（默认：3）
  --paugs PAUGS         对于伪选项：数据增强级别（默认：3）
  --lamb LAMB           对于伪选项：lambda 参数（默认：25）
  --cpu                 使用 CPU 版本（较慢）（默认：关闭）
  --arch ARCH           使用 AlexNet 或 ResNet（默认：AlexNet）
  --archspec {big,small}
                        AlexNet 的变体（默认：big）
  --ncl NCL             每个头的聚类数量（默认：3000）
  --hc HC               头的数量（默认：1）
  --device DEVICE       用于存储和模型计算的 GPU 设备
  --modeldevice MODELDEVICE
                        CNN 在其上运行的 GPU 编号
  --exp EXP             实验目录路径
  --workers WORKERS     工作线程数（默认：6）
  --imagenet-path IMAGENET_PATH
                        包含 `train` 和 `val` 文件夹的路径
  --comment COMMENT     tensorboardX 的名称
  --log-intv LOG_INTV   每隔 x 轮保存一次日志（默认：1）
  --log-iter LOG_ITER   每隔 x 批次记录一次日志（默认：200）

评估

线性评估

我们在该仓库中提供了线性评估的方法。只需通过 . ./scripts/download_models.sh 下载模型，然后运行 scripts/eval-alexnet.sh 或 scripts/eval-resnet.sh 即可。

Pascal VOC

我们遵循自监督视觉表征学习的标准评估协议。

• 对于分类任务：我们采用 DeepCluster 的 PyTorch 实现，并冻结 BatchNorm 层。
• 对于分割任务：我们使用 Colorization 论文 中的实现，该实现基于 FCN 仓库。注意：需要 Caffe 框架。
• 对于检测任务：我们遵循 Krähenbühl 等人 基于 Faster RCNN 的实现。注意：需要 Caffe 框架。

我们的提取的伪标签

如论文所示，我们从训练中生成的伪标签可用于快速训练具有常规交叉熵损失的神经网络。此外，这些标签似乎能够正确地将相似的图像归为一类。因此，我们提供这些标签供所有人使用。

AlexNet

你可以从我们最好的（原始）AlexNet 模型下载伪标签，该模型使用 10 组 3000 个聚类，下载链接为：这里。

ResNet

你可以从我们最好的 ResNet 模型下载伪标签，该模型使用 10 组 3000 个聚类，下载链接为：这里。

训练好的模型

你也可以通过运行以下命令下载我们训练好的模型：

$./scripts/download_models.sh

使用方法如下：

import torch
import models
d = torch.load('self-label_models/resnet-10x3k.pth')
m = models.resnet(num_classes = [3000]*10)
m.load_state_dict(d)

d = torch.load('self-label_models/alexnet-10x3k-wRot.pth')
m = models.alexnet(num_classes = [3000]*10)
m.load_state_dict(d)

参考文献

如果你使用了这段代码等，请引用以下论文：

Yuki M. Asano, Christian Rupprecht 和 Andrea Vedaldi. “通过同时进行聚类和表征学习实现自标签。” ICLR 会议论文集（2020 年）

@inproceedings{asano2020self,
  title={Self-labelling via simultaneous clustering and representation learning},
  author={Asano, Yuki M. and Rupprecht, Christian and Vedaldi, Andrea},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2020},
}

Self-Label 快速上手指南

Self-Label 是 ICLR 2020 论文《Self-labelling via simultaneous clustering and representation learning》的官方实现。该工具通过同时进行的聚类和表示学习，在无监督条件下为图像数据生成伪标签，从而训练出具有竞争力的图像表征模型（如 AlexNet, ResNet）。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐，需支持 CUDA)
• Python: > 3.6
• 深度学习框架: PyTorch > 1.0
• 硬件: NVIDIA GPU 及对应的 CUDA 环境（CPU 模式显著较慢，仅用于调试）
• 核心依赖:
  • NumPy
  • SciPy
  • TensorboardX (可选，用于可视化日志)

建议创建独立的虚拟环境并安装依赖：

pip install torch torchvision numpy scipy tensorboardx
# 或者使用 requirements.txt
pip install -r requirements.txt

提示：国内开发者可使用清华或阿里镜像源加速安装： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch torchvision numpy scipy tensorboardx

安装步骤

本项目无需复杂的编译安装过程，直接克隆仓库即可使用。

1. 克隆代码库

   git clone https://github.com/valeoai/self-label.git
   cd self-label
   

2. 下载预训练模型（可选） 如果您想直接评估或使用作者提供的预训练权重（包含标准的 torchvision ResNet 骨干网络），可以运行以下脚本：

   ./scripts/download_models.sh
   

3. 准备数据集 代码期望的数据集格式与 ImageNet 标准格式一致，即根目录下包含 train 和 val 文件夹。

   • 示例结构：

     /path/to/imagenet/
         ├── train/
         │   ├── n01440764/
         │   └── ...
         └── val/
             ├── n01440764/
             └── ...
     

基本使用

1. 无监督训练

您可以使用提供的脚本快速启动训练。请根据您的硬件资源修改脚本中的数据集路径。

训练 AlexNet:

./scripts/alexnet.sh

训练 ResNet-50:

./scripts/resnet.sh

注意：

• 您需要在脚本中指定 --imagenet-path 参数指向您的数据集目录。
• 训练需要足够的 GPU 显存来存储激活值。如果显存不足，可添加 --cpu 参数（速度会非常慢）。
• 默认配置针对多 GPU 环境优化。

自定义训练参数示例： 若需手动调整参数（如聚类数量、架构等），可直接调用 main.py：

python main.py \
  --arch resnet \
  --ncl 3000 \
  --epochs 280 \
  --batch-size 256 \
  --imagenet-path /path/to/your/dataset \
  --exp ./experiments/my_self_label_run

2. 加载预训练模型进行推理

如果您下载了预训练模型，可以使用标准的 PyTorch 方式加载。以下是加载新版 ResNet-50 模型的示例：

import torch
import torchvision

# 初始化模型 (注意 num_classes 需与预训练时的聚类数一致，此处为 3000)
model = torchvision.models.resnet50(pretrained=False, num_classes=3000)

# 加载权重
ckpt = torch.load('resnet50-10x3k_pp.pth') # 确保文件在当前目录或提供完整路径
model.load_state_dict(ckpt['state_dict'])

# 获取对应的伪标签 (如果需要)
pseudolabels = ckpt['labels']

# 将模型设置为评估模式
model.eval()

3. 线性评估 (Linear Evaluation)

训练完成后，通常需要进行线性评估以验证表征质量。

1. 确保已下载模型（参考“安装步骤”）。
2. 运行评估脚本：

   # 评估 AlexNet
   ./scripts/eval-alexnet.sh
   
   # 评估 ResNet
   ./scripts/eval-resnet.sh
   

4. 使用生成的伪标签

Self-Label 生成的伪标签可用于通过常规交叉熵损失快速微调神经网络。您可以直接从作者提供的链接下载 CSV 格式的标签文件（AlexNet 或 ResNet 版本），并在您的自定义训练循环中作为 Ground Truth 使用。