SupContrast: 监督对比学习

本仓库包含了以下论文在PyTorch中的参考实现，使用CIFAR作为说明性示例： (1) 监督对比学习。论文 (2) 视觉表示对比学习的简单框架。论文

更新

${\color{red}注意}$：如果您发现本仓库中的supcon损失实现难以解析，我们已经为您准备好了。Supcon损失本质上只是一个交叉熵损失（参见StableRep论文中的等式4）。因此，我们在这里提供了一个更清晰简洁的实现。希望这能帮到您。

ImageNet模型（使用动量编码器技巧的小批量大小）已在此处发布。它达到了超过79%的top-1准确率。

损失函数

losses.py中的损失函数SupConLoss接受特征（L2归一化）和标签作为输入，并返回损失。如果标签为None或未传递给它，则退化为SimCLR。

用法：

from losses import SupConLoss

# 定义具有温度参数`temp`的损失
criterion = SupConLoss(temperature=temp)

# features: [bsz, n_views, f_dim]
# `n_views`是每张图像的裁剪数量
# 最好在f_dim维度上进行L2归一化
features = ...
# labels: [bsz]
labels = ...

# SupContrast
loss = criterion(features, labels)
# 或 SimCLR
loss = criterion(features)
...

比较

CIFAR-10的结果：

CIFAR-100的结果：

ImageNet的结果（敬请期待）：

运行

您可以使用CUDA_VISIBLE_DEVICES设置适当数量的GPU，和/或通过--dataset cifar100切换到CIFAR100。 (1) 标准交叉熵

python main_ce.py --batch_size 1024 \
  --learning_rate 0.8 \
  --cosine --syncBN \

(2) 监督对比学习 预训练阶段：

python main_supcon.py --batch_size 1024 \
  --learning_rate 0.5 \
  --temp 0.1 \
  --cosine

您也可以指定--syncBN，但我发现它对SupContrast并不关键（syncBN 95.9% vs BN 96.0%）。

警告：目前，--syncBN没有效果，因为代码使用的是DataParallel而不是DistributedDataParaleel

线性评估阶段：

python main_linear.py --batch_size 512 \
  --learning_rate 5 \
  --ckpt /path/to/model.pth

(3) SimCLR 预训练阶段：

python main_supcon.py --batch_size 1024 \
  --learning_rate 0.5 \
  --temp 0.5 \
  --cosine --syncBN \
  --method SimCLR

--method SimCLR标志只是阻止将标签传递给SupConLoss准则。 线性评估阶段：

python main_linear.py --batch_size 512 \
  --learning_rate 1 \
  --ckpt /path/to/model.pth

在自定义数据集上：

python main_supcon.py --batch_size 1024 \
  --learning_rate 0.5  \ 
  --temp 0.1 --cosine \
  --dataset path \
  --data_folder ./path \
  --mean "(0.4914, 0.4822, 0.4465)" \
  --std "(0.2675, 0.2565, 0.2761)" \
  --method SimCLR

--data_folder必须采用./path/label/xxx.png的形式，遵循https://pytorch.org/docs/stable/torchvision/datasets.html#torchvision.datasets.ImageFolder 约定。

t-SNE可视化

(1) 标准交叉熵

(2) 监督对比学习

(3) SimCLR

参考文献

@Article{khosla2020supervised,
    title   = {Supervised Contrastive Learning},
    author  = {Prannay Khosla and Piotr Teterwak and Chen Wang and Aaron Sarna and Yonglong Tian and Phillip Isola and Aaron Maschinot and Ce Liu and Dilip Krishnan},
    journal = {arXiv preprint arXiv:2004.11362},
    year    = {2020},
}