SupContrast：有监督对比学习

本仓库提供了以下论文的 PyTorch 参考实现，并以 CIFAR 数据集作为示例： (1) 有监督对比学习。论文
(2) 视觉表征对比学习的简单框架。论文

更新

${\color{red}注意}$：如果您觉得解析本仓库中的 supcon 损失实现有些困难，我们为您准备了更简洁的版本。SupCon 损失本质上就是交叉熵损失（参见 StableRep 论文中的公式 4）。因此，我们在 这里 提供了一个更加清晰、简单的实现。希望对您有所帮助。

使用动量编码器技巧的小批量 ImageNet 模型已发布在 这里，其 top-1 准确率超过 79%。

损失函数

losses.py 中的 SupConLoss 损失函数接受 L2 归一化的 features 和 labels 作为输入，并返回损失值。如果 labels 为 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 上的结果：

	架构	设置	损失	准确率(%)
SupCrossEntropy	ResNet50	有监督	交叉熵	95.0
SupContrast	ResNet50	有监督	对比学习	96.0
SimCLR	ResNet50	无监督	对比学习	93.6

CIFAR-100 上的结果：

	架构	设置	损失	准确率(%)
SupCrossEntropy	ResNet50	有监督	交叉熵	75.3
SupContrast	ResNet50	有监督	对比学习	76.5
SimCLR	ResNet50	无监督	对比学习	70.7

ImageNet 上的结果（敬请期待）：

	架构	设置	损失	准确率(%)
SupCrossEntropy	ResNet50	有监督	交叉熵	-
SupContrast	ResNet50	有监督	对比学习	79.1（MoCo 技巧）
SimCLR	ResNet50	无监督	对比学习	-

运行

您可以使用 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% v.s. BN 96.0%）。

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

线性评估阶段：

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 标志的作用是阻止 labels 被传递给 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 必须遵循 https://pytorch.org/docs/stable/torchvision/datasets.html#torchvision.datasets.ImageFolder 的约定，格式为 ./path/label/xxx.png。

以及

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},
}

SupContrast 快速上手指南

SupContrast 是一个基于 PyTorch 的开源项目，实现了监督对比学习（Supervised Contrastive Learning）。相比传统的交叉熵损失函数，该方法在图像分类任务中通常能获得更高的准确率；若不提供标签，它也可退化为无监督的 SimCLR 方法。

环境准备

本项目依赖 Python 和 PyTorch 环境。建议在使用前确保系统已安装以下基础组件：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.6+
• PyTorch: 1.4+ (需包含 torchvision)
• GPU: 推荐使用 NVIDIA GPU 以加速训练（代码示例中使用了多卡批处理）

前置依赖安装：

pip install torch torchvision tensorboard numpy

提示：国内用户建议使用清华或阿里镜像源加速安装： pip install torch torchvision tensorboard numpy -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

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

1. 克隆仓库

   git clone https://github.com/HobbitLong/SupContrast.git
   cd SupContrast
   

2. 验证文件结构 确保目录下包含 losses.py（核心损失函数）、main_supcon.py（训练脚本）等关键文件。

基本使用

1. 核心损失函数调用

这是最精简的使用方式，可直接集成到你现有的 PyTorch 训练循环中。SupConLoss 接受归一化的特征向量和标签。

from losses import SupConLoss

# 定义损失函数，设置温度参数 temp
criterion = SupConLoss(temperature=0.1)

# 模拟输入数据
# features: [batch_size, n_views, feature_dim] 
# n_views 表示每张图像的增强视图数量（通常为 2）
# 建议在 feature_dim 维度进行 L2 归一化
features = ... 
# labels: [batch_size]
# 如果传入 labels，则为 Supervised Contrastive Learning
# 如果不传 labels (或设为 None)，则自动退化为 SimCLR
labels = ...

# 计算损失
loss = criterion(features, labels)

2. 运行训练示例

项目提供了完整的训练脚本，支持 CIFAR-10/100 数据集。

场景 A：监督对比学习 (Supervised Contrastive Learning) 使用标签进行预训练，通常能获得最佳效果。

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

场景 B：无监督对比学习 (SimCLR) 不传入标签，复现 SimCLR 算法。

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

场景 C：线性评估 (Linear Evaluation) 在预训练模型基础上冻结骨干网络，仅训练分类头进行评估。

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

3. 自定义数据集

若要使用自己的图像数据集，需遵循 PyTorch ImageFolder 的目录结构（即 ./data_folder/label_name/image.png），并指定均值和标准差。

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

注意：多卡训练时，可通过 CUDA_VISIBLE_DEVICES 环境变量指定 GPU 编号。例如：CUDA_VISIBLE_DEVICES=0,1 python main_supcon.py ...