第14章 自监督学习¶

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimCLR(nn.Module):  # 继承nn.Module定义网络层
    def __init__(self, encoder, projection_dim=128):
        super(SimCLR, self).__init__()
        self.encoder = encoder
        self.projection = nn.Sequential(
            nn.Linear(encoder.output_dim, encoder.output_dim),
            nn.ReLU(),
            nn.Linear(encoder.output_dim, projection_dim)
        )

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

def nt_xent_loss(z_i, z_j, temperature=0.5):
    """Normalized Temperature-scaled Cross Entropy Loss

    注意：必须屏蔽对角线（自相似度=1.0），否则模型会
    "作弊"匹配自身而非正样本对，导致训练无效。
    """
    batch_size = z_i.shape[0]
    N = 2 * batch_size

    # 拼接正负样本
    z = torch.cat([z_i, z_j], dim=0)  # torch.cat沿已有维度拼接张量

    # 计算相似度矩阵 (2N x 2N)
    sim = F.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=2)  # unsqueeze增加一个维度  # F.xxx PyTorch函数式API

    # ⚠️ 屏蔽对角线（自相似度），设为 -inf 使其在 softmax 中贡献为 0
    mask = torch.eye(N, dtype=torch.bool, device=z_i.device)
    sim = sim.masked_fill(mask, float('-inf'))

    # 创建标签：z_i[k] 的正样本是 z_j[k]（索引 batch_size+k）
    # 屏蔽对角线后标签索引需要调整（每行少了一个自身元素）
    # 但 cross_entropy 会忽略 -inf 位置，标签仍指向原始列索引
    labels = torch.arange(batch_size, device=z_i.device)
    labels = torch.cat([labels + batch_size, labels], dim=0)

    # 温度缩放
    sim = sim / temperature

    # 损失
    loss = F.cross_entropy(sim, labels)  # F.cross_entropy PyTorch函数式交叉熵损失

    return loss

import copy

class MoCo(nn.Module):
    def __init__(self, encoder, K=65536, m=0.999, temperature=0.07):
        super(MoCo, self).__init__()
        self.K = K
        self.m = m
        self.T = temperature

        # Query编码器
        self.encoder_q = encoder

        # Key编码器（动量更新）— 必须深拷贝，不能共享引用
        self.encoder_k = copy.deepcopy(encoder)
        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):  # zip按位置配对
            param_k.data.copy_(param_q.data)
            param_k.requires_grad = False

        # 队列
        self.register_buffer("queue", torch.randn(encoder.output_dim, K))
        self.queue = F.normalize(self.queue, dim=0)
        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

    @torch.no_grad()  # 禁用梯度计算，节省内存
    def _momentum_update_key_encoder(self):
        """动量更新key编码器"""
        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys):
        """更新队列"""
        batch_size = keys.shape[0]

        ptr = int(self.queue_ptr)

        # 更新队列
        self.queue[:, ptr:ptr + batch_size] = keys.T

        # 更新指针
        ptr = (ptr + batch_size) % self.K
        self.queue_ptr[0] = ptr

    def forward(self, im_q, im_k):
        # Query
        q = self.encoder_q(im_q)
        q = F.normalize(q, dim=1)

        # Key（不计算梯度）
        with torch.no_grad():
            self._momentum_update_key_encoder()
            k = self.encoder_k(im_k)
            k = F.normalize(k, dim=1)

        # 计算相似度
        l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
        l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])  # 分离计算图，不参与梯度计算

        logits = torch.cat([l_pos, l_neg], dim=1)
        logits /= self.T

        # 标签（正样本在第一个位置）
        labels = torch.zeros(logits.shape[0], dtype=torch.long).to(q.device)

        loss = F.cross_entropy(logits, labels)

        # 更新队列
        self._dequeue_and_enqueue(k)

        return loss

class MAE(nn.Module):
    def __init__(self, encoder, decoder, mask_ratio=0.75):
        super(MAE, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.mask_ratio = mask_ratio

    def random_masking(self, x, mask_ratio):
        """随机掩码"""
        N, L, D = x.shape
        len_keep = int(L * (1 - mask_ratio))

        # 随机噪声
        noise = torch.rand(N, L, device=x.device)

        # 选择保留的patch
        ids_shuffle = torch.argsort(noise, dim=1)
        ids_restore = torch.argsort(ids_shuffle, dim=1)

        ids_keep = ids_shuffle[:, :len_keep]
        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))

        # 生成mask
        mask = torch.ones([N, L], device=x.device)
        mask[:, :len_keep] = 0
        mask = torch.gather(mask, dim=1, index=ids_restore)

        return x_masked, mask, ids_restore

    def forward(self, x):
        # 随机掩码
        x_masked, mask, ids_restore = self.random_masking(x, self.mask_ratio)

        # 编码
        latent = self.encoder(x_masked)

        # 解码
        pred = self.decoder(latent, ids_restore)

        # 损失（只计算被掩码的部分）
        loss = (pred - x) ** 2
        loss = loss.mean(dim=-1)
        loss = (loss * mask).sum() / mask.sum() if mask.sum() > 0 else loss.sum()

        return loss, pred, mask

import torch.optim as optim
from torchvision import datasets, transforms

# 数据增强
class SimCLRTransform:
    def __init__(self, size=224):
        self.transform = transforms.Compose([
            transforms.RandomResizedCrop(size),
            transforms.RandomHorizontalFlip(),
            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([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    def __call__(self, x):  # __call__使实例可像函数一样调用
        return self.transform(x), self.transform(x)

# 数据集
transform = SimCLRTransform()
dataset = datasets.CIFAR10('./data', train=True, download=True, transform=transform)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=256, shuffle=True)  # DataLoader批量加载数据

# 模型
encoder = ResNet18()
model = SimCLR(encoder).cuda()

# 优化器
optimizer = optim.Adam(model.parameters(), lr=3e-4)

# 训练
def train_simclr(model, dataloader, epochs=100):
    model.train()  # train()训练模式
    for epoch in range(epochs):
        total_loss = 0.0
        for (x_i, x_j), _ in dataloader:
            x_i, x_j = x_i.cuda(), x_j.cuda()

            # 前向传播
            z_i = model(x_i)
            z_j = model(x_j)

            # 损失
            loss = nt_xent_loss(z_i, z_j)

            # 反向传播
            optimizer.zero_grad()  # 清零梯度
            loss.backward()  # 反向传播计算梯度
            optimizer.step()  # 更新参数

            total_loss += loss.item()  # 将单元素张量转为Python数值

        print(f'Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}')

# 训练
train_simclr(model, dataloader, epochs=100)