SimCLR pytorch 代码

SimCLR (Simple Framework for Contrastive Learning of Visual Representations) 是一种无监督的视觉表示学习方法，它通过对比样本对（通常是在数据增强后的版本）来训练深度神经网络。在PyTorch中实现SimCLR的基本步骤包括：

1. **数据预处理**：加载图像数据并应用随机数据增强，如翻转、裁剪和颜色调整，创建正负样本对。

import torch
from torchvision.transforms import transforms

transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

def get_augmented_pairs(x):
    x_pos = transform(x)
    x_neg = transform(x)
    return x_pos, x_neg

2. **构建模型**：使用PyTorch构建卷积神经网络（CNN），例如ResNet、ViT等作为基础模型。

import torch.nn as nn
import torchvision.models as models

model = models.resnet50(pretrained=False)
# 去掉分类层
num_features = model.fc.in_features
model.fc = nn.Identity()

3. **构建SimCLR架构**：这里通常包含一个投影头（projector），用于将来自原始模型的特征映射到更低维度的空间，并一个加权平均池化层（weight tying trick）。

class ProjectionHead(nn.Module):
    def __init__(self, num_features, projection_dim):
        super().__init__()
        self.projector = nn.Sequential(
            nn.Linear(num_features, projection_dim),
            nn.ReLU(inplace=True),
            nn.Linear(projection_dim, projection_dim),
        )

    def forward(self, x):
        return self.projector(x)

projection_head = ProjectionHead(num_features, proj_dim=128)

4. **优化器设置**：使用AdamW或其他优化器，并设置学习率衰减策略（如warm-up）。

5. **训练循环**：对于每个批次的数据，计算正负样本对的嵌入之间的相似度损失，然后更新模型参数。

optimizer = torch.optim.AdamW(model.parameters(), lr=0.005)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=1000)

for epoch in range(num_epochs):
    # 训练过程...
    pos_embeddings, neg_embeddings = ... # 获取一对样本的正负嵌入
    loss = ... # 使用某种对比性损失函数（如NT-Xent）
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    scheduler.step()