📚 课程目录第3章 SimSiam孪生网络实现

3. SimSiam孪生网络实现

任务目标

实现SimSiam孪生网络，包含停止梯度（stop-gradient）操作与预测器（predictor）模块，在STL-10数据集上完成训练，并通过对比实验验证该机制能有效避免表示崩溃（即所有样本被映射到同一常数向量）。

核心实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# ---------- SimSiam 基础模块 ----------
class Projector(nn.Module):
    """将编码器输出映射到低维空间"""
    def __init__(self, in_dim=512, hidden_dim=2048, out_dim=2048):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, out_dim),
            nn.BatchNorm1d(out_dim)
        )

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

class Predictor(nn.Module):
    """SimSiam 特有的预测器，仅在在线分支使用"""
    def __init__(self, in_dim=2048, hidden_dim=512, out_dim=2048):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, out_dim)
        )

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

class SimSiam(nn.Module):
    def __init__(self, base_encoder=torchvision.models.resnet18):
        super().__init__()
        # 编码器：去掉原始 ResNet 的最后的全连接层
        self.encoder = base_encoder(num_classes=0)
        feat_dim = self.encoder.fc.in_features if hasattr(self.encoder, 'fc') else 512
        self.projector = Projector(in_dim=feat_dim)
        self.predictor = Predictor()

    def forward(self, x1, x2):
        # 两个增强视图分别经过编码器+投影器
        z1 = self.projector(self.encoder(x1))
        z2 = self.projector(self.encoder(x2))

        # 在线分支：预测器 + 停止梯度
        p1 = self.predictor(z1)
        p2 = self.predictor(z2)

        # 目标分支：停止梯度（不更新该路径的梯度）
        z1_detached = z1.detach()
        z2_detached = z2.detach()

        return p1, p2, z1_detached, z2_detached

    def loss_fn(self, p1, p2, z1, z2):
        """负余弦相似度损失"""
        def neg_cos_sim(p, z):
            p = F.normalize(p, dim=1)
            z = F.normalize(z, dim=1)
            return -(p * z).sum(dim=1).mean()

        return (neg_cos_sim(p1, z2) + neg_cos_sim(p2, z1)) / 2

# ---------- 数据增强（STL-10 专用） ----------
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(96, scale=(0.2, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.4, 0.4, 0.4, 0.1),
    transforms.RandomGrayscale(p=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# ---------- 训练配置 ----------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
batch_size = 64
epochs = 30
lr = 0.03

# 加载 STL-10 无标签训练集（unlabeled 包含 10 万张图像）
train_dataset = torchvision.datasets.STL10(
    root='./data', split='unlabeled', download=True, transform=train_transform
)
train_loader = DataLoader(train_dataset, batch_size=batch_size,
                          shuffle=True, num_workers=4, drop_last=True)

model = SimSiam().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=lr,
                            momentum=0.9, weight_decay=5e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

# ---------- 训练循环 ----------
model.train()
for epoch in range(epochs):
    total_loss = 0.0
    for batch_idx, (images, _) in enumerate(train_loader):
        # 同一批次生成两个不同增强视图
        x1, x2 = images.to(device), images.to(device)

        p1, p2, z1_det, z2_det = model(x1, x2)
        loss = model.loss_fn(p1, p2, z1_det, z2_det)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    scheduler.step()
    avg_loss = total_loss / len(train_loader)
    print(f'Epoch [{epoch+1}/{epochs}]  Loss: {avg_loss:.4f}')

# ---------- 崩溃检测：检查表示方差 ----------
model.eval()
with torch.no_grad():
    # 取一个批次，检查编码器输出的标准差
    sample_images, _ = next(iter(train_loader))
    sample_images = sample_images.to(device)
    z = model.projector(model.encoder(sample_images))
    std_per_dim = z.std(dim=0)
    mean_std = std_per_dim.mean().item()
    print(f'表示标准差均值: {mean_std:.4f} (若接近0则发生崩溃)')

运行与验证

环境依赖：torch、torchvision、numpy。首次运行会自动下载STL-10数据集（约2.5GB），需确保网络连接。

运行方式：直接执行上述脚本。训练约需30分钟（单张V100 GPU），若资源有限可将epochs调至5进行快速验证。

通过判定标准：

训练损失应持续下降并稳定在-0.8至-0.95区间（负余弦相似度的理论下界为-1）；
训练结束后，崩溃检测输出的表示标准差均值应大于0.1。若该值接近0（如<0.01），说明所有样本被映射到同一向量，模型发生崩溃；
作为对照，若移除detach()操作（即不使用停止梯度），相同训练配置下标准差均值会显著下降至接近0，从而验证停止梯度是SimSiam避免崩溃的关键机制。

关键说明

本实现使用ResNet-18作为编码器，输入图像为STL-10原始尺寸96×96。停止梯度通过detach()实现，确保目标分支的梯度不反向传播；预测器仅作用于在线分支，这是SimSiam与SimCLR等方法的本质区别。若需验证崩溃避免机制的有效性，可对比删除detach()后的训练结果，观察表示方差是否显著降低。