AI自监督学习折腾手记

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms
from torch.utils.data import DataLoader, Dataset
from torchvision.datasets import ImageFolder

# 检查 CUDA 可用性
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

from torchvision.transforms import transforms

def get_simclr_augmentation(size=224):
    return transforms.Compose([
        transforms.RandomResizedCrop(size, scale=(0.2, 1.0)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomApply([transforms.ColorJitter(0.8, 0.8, 0.8, 0.2)], p=0.8),
        transforms.RandomGrayscale(p=0.2),
        transforms.GaussianBlur(kernel_size=int(0.1 * size), sigma=(0.1, 2.0)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

class SimCLRDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.dataset = ImageFolder(root_dir)
        self.transform = transform

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        img, _ = self.dataset[idx]
        if self.transform:
            img1 = self.transform(img)
            img2 = self.transform(img)
        else:
            img1 = img2 = img
        return img1, img2

# 数据加载器
train_dataset = SimCLRDataset(
    '/path/to/unlabeled_data',
    transform=get_simclr_augmentation()
)

train_loader = DataLoader(
    train_dataset,
    batch_size=256,  # 每张卡
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    drop_last=True
)

from torchvision.models import resnet50

class SimCLR(nn.Module):
    def __init__(self, base_model='resnet50', out_dim=128):
        super(SimCLR, self).__init__()
        self.encoder = resnet50(weights=None)

        # 移除最后的全连接层
        self.encoder.fc = nn.Identity()

        # Projection head: 2048 -> 512 -> 128
        self.projection_head = nn.Sequential(
            nn.Linear(2048, 512),
            nn.ReLU(),
            nn.Linear(512, out_dim)
        )

    def forward(self, x):
        h = self.encoder(x)
        z = self.projection_head(h)
        return h, z

model = SimCLR().to(device)

import torch.nn.functional as F

def info_nce_loss(features, temperature=0.5):
    batch_size = features.shape[0]
    labels = torch.arange(batch_size).to(features.device)

    # 计算相似度矩阵
    similarity_matrix = F.cosine_similarity(
        features.unsqueeze(1), features.unsqueeze(0), dim=2
    )

    # 去掉对角线（自己和自己不用比）
    mask = torch.eye(batch_size, dtype=torch.bool).to(features.device)
    similarity_matrix = similarity_matrix[~mask].view(batch_size, -1)

    # 计算损失
    loss = F.cross_entropy(similarity_matrix / temperature, labels)
    return loss

from torch.cuda.amp import autocast, GradScaler

# 优化器和学习率调度
optimizer = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=100, eta_min=1e-6
)

scaler = GradScaler()

num_epochs = 100
temperature = 0.5
accumulation_steps = 4  # 模拟更大的 batch size

model.train()
for epoch in range(num_epochs):
    epoch_loss = 0.0
    optimizer.zero_grad()

    for batch_idx, (img1, img2) in enumerate(train_loader):
        img1, img2 = img1.to(device), img2.to(device)

        # 前向传播
        with autocast():
            _, z1 = model(img1)
            _, z2 = model(img2)

            # 合并两个视图
            z = torch.cat([z1, z2], dim=0)
            loss = info_nce_loss(z, temperature)

        # 反向传播（考虑梯度累积）
        scaler.scale(loss / accumulation_steps).backward()

        if (batch_idx + 1) % accumulation_steps == 0:
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

        epoch_loss += loss.item()

    scheduler.step()
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss/len(train_loader):.4f}")

    # 定期保存检查点
    if (epoch + 1) % 10 == 0:
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict(),
        }, f'simclr_checkpoint_epoch_{epoch+1}.pth')

# 梯度累积示例
accumulation_steps = 8
for i, batch in enumerate(train_loader):
    loss = model(batch)
    scaler.scale(loss / accumulation_steps).backward()

    if (i + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

# 调整后的增强策略
transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.4, 1.0)),  # 调高下限
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomApply([transforms.ColorJitter(0.6, 0.6, 0.6, 0.1)], p=0.6),  # 减弱颜色抖动
    transforms.RandomGrayscale(p=0.1),  # 降低灰度概率
    transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 1.5)),  # 适度模糊
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 温度调度：开始高一点，后期降低
def get_temperature(epoch, total_epochs):
    start_temp = 0.1
    end_temp = 0.07
    return start_temp + (end_temp - start_temp) * epoch / total_epochs

from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

def visualize_features(model, data_loader, device):
    model.eval()
    features = []
    labels = []

    with torch.no_grad():
        for imgs, lbls in data_loader:
            imgs = imgs.to(device)
            h, _ = model(imgs)
            features.append(h.cpu().numpy())
            labels.append(lbls.numpy())

    features = np.concatenate(features, axis=0)
    labels = np.concatenate(labels, axis=0)

    # t-SNE 降维
    tsne = TSNE(n_components=2, random_state=42)
    features_2d = tsne.fit_transform(features)

    # 可视化
    plt.figure(figsize=(10, 8))
    scatter = plt.scatter(features_2d[:, 0], features_2d[:, 1], c=labels, cmap='tab10')
    plt.colorbar(scatter)
    plt.title('Feature Space Visualization')
    plt.savefig('feature_space.png')