SimCLR

由T. Chen等人提出的用于视觉表征对比学习的简单框架SimCLR的PyTorch实现。 支持以下功能：

• 分布式数据并行训练
• 全局批归一化
• LARS（逐层自适应速率缩放）优化器。

论文链接

在Google Colab Notebook中打开SimCLR（支持TPU）

在tensorboard.dev上查看SimCLR结果对比：

快速入门（微调线性分类器）

此步骤将下载一个预训练模型并训练线性分类器，该分类器在STL-10测试集上的准确率应达到±82.9%。

git clone https://github.com/spijkervet/SimCLR.git && cd SimCLR
wget https://github.com/Spijkervet/SimCLR/releases/download/1.2/checkpoint_100.tar
sh setup.sh || python3 -m pip install -r requirements.txt || exit 1
conda activate simclr
python linear_evaluation.py --dataset=STL10 --model_path=. --epoch_num=100 --resnet resnet50

CPU

wget https://github.com/Spijkervet/SimCLR/releases/download/1.1/checkpoint_100.tar -O checkpoint_100.tar
python linear_evaluation.py --model_path=. --epoch_num=100 --resnet=resnet18 --logistic_batch_size=32

simclr软件包

SimCLR现已作为Python包提供！只需运行即可在您的项目中使用：

pip install simclr

然后您可以直接导入SimCLR：

from simclr import SimCLR

encoder = ResNet(...)
projection_dim = 64
n_features = encoder.fc.in_features  # 获取最后一层全连接层的维度
model = SimCLR(encoder, projection_dim, n_features)

训练ResNet编码器：

只需运行以下命令即可使用SimCLR在CIFAR-10数据集上预训练一个ResNet编码器：

python main.py --dataset CIFAR10

分布式训练

通过分布式数据并行（DDP）训练：

CUDA_VISIBLE_DEVICES=0 python main.py --nodes 2 --nr 0
CUDA_VISIBLE_DEVICES=1 python main.py --nodes 2 --nr 1
CUDA_VISIBLE_DEVICES=2 python main.py --nodes 2 --nr 2
CUDA_VISIBLE_DEVICES=N python main.py --nodes 2 --nr 3

结果

以下是基于SimCLR学习到的（冻结）表征训练的线性分类器的Top-1准确率：

方法	批量大小	ResNet	投影输出维度	轮数	优化器	STL-10	CIFAR-10
SimCLR + 线性评估	256	ResNet50	64	100	Adam	0.829	0.833
SimCLR + 线性评估	256	ResNet50	64	100	LARS	0.783	-
SimCLR + 线性评估	256	ResNet18	64	100	Adam	0.765	-
SimCLR + 线性评估	256	ResNet18	64	40	Adam	0.719	-
SimCLR + 线性评估	512	ResNet18	64	40	Adam	0.71	-
逻辑回归	-	-	-	40	Adam	0.358	0.389

预训练模型

ResNet（批量大小，轮数）	优化器	STL-10 Top-1
ResNet50 (256, 100)	Adam	0.829
ResNet18 (256, 100)	Adam	0.765
ResNet18 (256, 40)	Adam	0.719

python linear_evaluation.py --model_path=. --epoch_num=100

LARS优化器

LARS优化器实现于modules/lars.py。可通过将config/config.yaml中的优化器设置调整为：optimizer: "LARS"来启用。该优化器目前仍处于实验阶段，尚未经过充分测试。

什么是SimCLR？

SimCLR是一种“用于视觉表征对比学习的简单框架”。对比预测任务定义在成对的增强样本上，每个小批量包含2N个样本。一张图像的两个增强版本被视为相关联的“正”样本对（x_i和x_j）。其余2(N - 1)个增强样本则被视为负样本。对比预测任务的目标是针对给定的x_i，在负样本集中识别出x_j。

使用方法

运行以下命令以设置conda环境：

sh setup.sh
conda activate simclr

或者使用pip：

pip install -r requirements.txt

然后，只需运行即可进行单GPU或CPU训练：

python main.py

对于分布式训练（DDP），请为每个节点中的每个进程使用以下命令，其中N为您希望分配给该进程的GPU编号：

CUDA_VISIBLE_DEVICES=0 python main.py --nodes 2 --nr 0
CUDA_VISIBLE_DEVICES=1 python main.py --nodes 2 --nr 1
CUDA_VISIBLE_DEVICES=2 python main.py --nodes 2 --nr 2
CUDA_VISIBLE_DEVICES=N python main.py --nodes 2 --nr 3

--nr对应于我们为训练提供的N个节点中的进程编号。

测试

要测试已训练的模型，请确保将config/config.yaml中的model_path变量设置为训练的日志ID（例如logs/0）。 将epoch_num设置为您希望加载检查点的轮次（例如40）。

python linear_evaluation.py

或就地执行：

python linear_evaluation.py --model_path=./save --epoch_num=40

配置

训练的配置可在config/config.yaml中找到。我个人更倾向于使用文件而不是长串的参数来配置运行。以下是一个config.yaml文件示例：

# 训练选项
batch_size: 256
workers: 16
start_epoch: 0
epochs: 40
dataset_dir: "./datasets"

# 模型选项
resnet: "resnet18"
normalize: True
projection_dim: 64

# 损失选项
temperature: 0.5

# 重载选项
model_path: "logs/0" # 设置为包含`checkpoint_##.tar`的目录
epoch_num: 40 # 设置为检查点编号

# 逻辑回归选项
logistic_batch_size: 256
logistic_epochs: 100

日志记录与TensorBoard

要在TensorBoard中查看结果，请运行：

tensorboard --logdir runs

优化器与学习率调度

本实现提供了Adam优化器和LARS优化器，并可选择使用余弦退火学习率调度来衰减学习率。优化器和权重衰减可以在config/config.yaml文件中进行配置。

依赖项

torch
torchvision
tensorboard
pyyaml

SimCLR 快速上手指南

SimCLR 是一个用于视觉表示对比学习的简单框架（PyTorch 实现）。本指南将帮助你快速配置环境、安装依赖并运行预训练模型评估或从头开始训练。

环境准备

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

• 操作系统: Linux 或 macOS (Windows 需通过 WSL 运行)
• Python: 3.6+
• 硬件: 支持 CUDA 的 GPU（推荐）或 CPU
• 前置依赖:
  • conda (推荐用于环境管理) 或 pip
  • Git

核心依赖库：

• torch
• torchvision
• tensorboard
• pyyaml

国内加速建议：建议使用清华源或阿里源加速 Python 包和 Conda 环境的安装，以提升下载速度。

安装步骤

你可以选择使用 Conda 创建隔离环境（推荐），或直接使用 pip 安装。

方法一：使用 Conda（推荐）

1. 克隆项目代码：

   git clone https://github.com/spijkervet/SimCLR.git && cd SimCLR
   

2. 创建并激活环境： 运行提供的设置脚本（该脚本会自动处理依赖安装）：

   sh setup.sh
   conda activate simclr
   

   如果 setup.sh 执行失败，可手动安装依赖：

   python3 -m pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   conda create -n simclr python=3.8
   conda activate simclr
   pip install torch torchvision tensorboard pyyaml -i https://pypi.tuna.tsinghua.edu.cn/simple
   

方法二：作为 Python 包安装

如果你只想在现有项目中调用 SimCLR 模块：

pip install simclr -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

场景 1：快速评估预训练模型（微调线性分类器）

这是最简单的入门方式。以下命令会下载一个在 STL-10 数据集上预训练的 ResNet50 模型，并训练一个线性分类器进行评估（预期准确率约为 82.9%）。

# 下载预训练权重
wget https://github.com/Spijkervet/SimCLR/releases/download/1.2/checkpoint_100.tar

# 运行线性评估
python linear_evaluation.py --dataset=STL10 --model_path=. --epoch_num=100 --resnet resnet50

CPU 用户示例（使用较小的 ResNet18 模型）：

wget https://github.com/Spijkervet/SimCLR/releases/download/1.1/checkpoint_100.tar -O checkpoint_100.tar
python linear_evaluation.py --model_path=. --epoch_num=100 --resnet=resnet18 --logistic_batch_size=32

场景 2：从头开始训练编码器

若要在 CIFAR-10 数据集上使用 SimCLR 框架预训练 ResNet 编码器：

python main.py --dataset CIFAR10

场景 3：在代码中直接调用

如果你已通过 pip install simclr 安装了包，可以在 Python 脚本中直接使用：

from simclr import SimCLR
from torchvision.models import resnet50

# 初始化编码器
encoder = resnet50()
projection_dim = 64
n_features = encoder.fc.in_features  # 获取最后一层全连接层的维度

# 构建 SimCLR 模型
model = SimCLR(encoder, projection_dim, n_features)

进阶：分布式训练 (DDP)

如果你有多张 GPU，可以使用以下命令启动分布式训练（假设使用 2 个节点）：

CUDA_VISIBLE_DEVICES=0 python main.py --nodes 2 --nr 0
CUDA_VISIBLE_DEVICES=1 python main.py --nodes 2 --nr 1

查看训练结果

训练完成后，使用 TensorBoard 查看损失曲线和准确率：

tensorboard --logdir runs

然后在浏览器中访问显示的地址（通常为 http://localhost:6006）。

配置提示：所有训练参数（如 batch size、epochs、优化器类型等）均可在 config/config.yaml 文件中修改。