📚 课程目录第2章 MoCo动量对比实现

2. MoCo动量对比实现

任务目标

实现MoCo v1（Momentum Contrast）的核心机制：动量编码器、负样本队列以及InfoNCE损失函数，并在ImageNet子集（10类，每类约130张训练图）上完成训练与评估，验证对比学习表征质量。

核心机制说明

MoCo通过维护一个大型负样本队列（而非依赖大batch size）来提供充足的负样本。查询编码器（query encoder）与动量编码器（momentum encoder）共享相同网络结构，但动量编码器的参数通过指数移动平均（EMA）更新，保证队列中负样本特征的一致性。InfoNCE损失将每个样本的增强视图（query）与对应正样本（key）拉近，同时推远队列中的负样本。

环境依赖

torch==2.0.1
torchvision==0.15.2
Pillow==10.0.0
numpy==1.24.3

完整实现代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as T
from torch.utils.data import DataLoader, Subset
import numpy as np
from PIL import Image
import os, random

# ---------- 1. 数据准备：ImageNet子集（10类） ----------
def prepare_subset(data_dir='/data/imagenet', num_classes=10, samples_per_class=130):
    """从ImageNet训练集中抽取前num_classes类，每类samples_per_class张"""
    full_set = torchvision.datasets.ImageFolder(
        root=os.path.join(data_dir, 'train'),
        transform=T.Compose([
            T.RandomResizedCrop(224, scale=(0.2, 1.0)),
            T.RandomHorizontalFlip(),
            T.ColorJitter(0.4, 0.4, 0.4, 0.1),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]))
    # 按类别筛选
    class_indices = {c: [] for c in range(num_classes)}
    for idx, (_, label) in enumerate(full_set.samples):
        if label < num_classes and len(class_indices[label]) < samples_per_class:
            class_indices[label].append(idx)
    selected = [i for c in class_indices.values() for i in c]
    return Subset(full_set, selected)

# ---------- 2. 编码器（ResNet-18，输出128维） ----------
class Encoder(nn.Module):
    def __init__(self, feat_dim=128):
        super().__init__()
        resnet = torchvision.models.resnet18(weights=None)
        self.backbone = nn.Sequential(*list(resnet.children())[:-1])
        self.fc = nn.Linear(512, feat_dim)
    def forward(self, x):
        feat = self.backbone(x).flatten(1)
        return F.normalize(self.fc(feat), dim=1)

# ---------- 3. MoCo核心实现 ----------
class MoCo(nn.Module):
    def __init__(self, dim=128, K=4096, m=0.999, T=0.07):
        super().__init__()
        self.K = K          # 队列大小
        self.m = m          # 动量系数
        self.T = T          # 温度
        self.encoder_q = Encoder(dim)
        self.encoder_k = Encoder(dim)
        # 初始化动量编码器参数
        for p_q, p_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            p_k.data.copy_(p_q.data)
            p_k.requires_grad = False
        # 负样本队列
        self.register_buffer('queue', F.normalize(torch.randn(K, dim), dim=1))
        self.register_buffer('queue_ptr', torch.zeros(1, dtype=torch.long))

    @torch.no_grad()
    def _momentum_update(self):
        for p_q, p_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            p_k.data = p_k.data * self.m + p_q.data * (1. - self.m)

    @torch.no_grad()
    def _dequeue_enqueue(self, keys):
        batch_size = keys.shape[0]
        ptr = int(self.queue_ptr)
        assert self.K % batch_size == 0
        self.queue[ptr:ptr+batch_size] = keys
        ptr = (ptr + batch_size) % self.K
        self.queue_ptr[0] = ptr

    def forward(self, im_q, im_k):
        # 查询特征
        q = self.encoder_q(im_q)
        # 动量编码器计算key（无梯度）
        with torch.no_grad():
            self._momentum_update()
            k = self.encoder_k(im_k)
        # InfoNCE损失
        l_pos = torch.einsum('nc,nc->n', q, k).unsqueeze(-1)  # [N,1]
        l_neg = torch.einsum('nc,ck->nk', q, self.queue.clone().detach())  # [N,K]
        logits = torch.cat([l_pos, l_neg], dim=1) / self.T
        labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()
        loss = F.cross_entropy(logits, labels)
        # 更新队列
        self._dequeue_enqueue(k)
        return loss

# ---------- 4. 训练与评估 ----------
def train_moco(model, train_loader, epochs=50, lr=0.03):
    optimizer = torch.optim.SGD(model.encoder_q.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 i, (im_q, im_k, _) in enumerate(train_loader):
            im_q, im_k = im_q.cuda(), im_k.cuda()
            loss = model(im_q, im_k)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        scheduler.step()
        if (epoch+1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{epochs}] Loss: {total_loss/len(train_loader):.4f}')

def evaluate(model, val_loader):
    """线性评估：冻结编码器，训练线性分类头"""
    model.eval()
    feat_dim = 128
    classifier = nn.Linear(feat_dim, 10).cuda()
    opt = torch.optim.Adam(classifier.parameters(), lr=1e-3)
    ce = nn.CrossEntropyLoss()
    # 提取特征
    feats, labels = [], []
    with torch.no_grad():
        for x, y in val_loader:
            x = x.cuda()
            feat = model.encoder_q(x)
            feats.append(feat.cpu())
            labels.append(y)
    feats = torch.cat(feats)
    labels = torch.cat(labels)
    # 训练线性层
    dataset = torch.utils.data.TensorDataset(feats, labels)
    loader = DataLoader(dataset, batch_size=256, shuffle=True)
    for epoch in range(30):
        for xb, yb in loader:
            xb, yb = xb.cuda(), yb.cuda()
            out = classifier(xb)
            loss = ce(out, yb)
            opt.zero_grad()
            loss.backward()
            opt.step()
    # 评估
    classifier.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for xb, yb in loader:
            xb, yb = xb.cuda(), yb.cuda()
            pred = classifier(xb).argmax(dim=1)
            correct += (pred == yb).sum().item()
            total += yb.size(0)
    return correct / total

if __name__ == '__main__':
    # 数据加载（需先运行prepare_subset生成索引）
    train_set = prepare_subset()
    # 为对比学习构造两个增强视图
    class TwoCropsTransform:
        def __init__(self, base_transform):
            self.base = base_transform
        def __call__(self, x):
            return self.base(x), self.base(x)
    train_loader = DataLoader(train_set, batch_size=256, shuffle=True,
                              num_workers=4, drop_last=True)
    # 注意：此处train_set的transform已包含增强，需替换为TwoCrops
    train_set.dataset.transform = TwoCropsTransform(train_set.dataset.transform)
    # 验证集（ImageNet val，10类）
    val_transform = T.Compose([
        T.Resize(256), T.CenterCrop(224),
        T.ToTensor(), T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
    ])
    val_set = torchvision.datasets.ImageFolder(
        root='/data/imagenet/val',
        transform=val_transform)
    val_indices = [i for i, (_, l) in enumerate(val_set.samples) if l < 10]
    val_loader = DataLoader(Subset(val_set, val_indices), batch_size=256, shuffle=False)

    model = MoCo(dim=128, K=4096, m=0.999, T=0.07).cuda()
    train_moco(model, train_loader, epochs=50)
    acc = evaluate(model, val_loader)
    print(f'Linear Evaluation Accuracy: {acc:.4f}')

运行方式

# 1. 准备数据（ImageNet格式目录结构）
#    /data/imagenet/train/ 包含类别子文件夹
#    /data/imagenet/val/   包含类别子文件夹
# 2. 运行训练
python moco_implementation.py

验证方法与通过标准

训练收敛验证：训练过程中InfoNCE损失应持续下降，50个epoch后损失通常低于0.5（具体值取决于数据子集难度）。
表征质量评估：通过线性评估协议（冻结编码器，训练线性分类头）在10类ImageNet子集上测试，若top-1准确率高于60%则证明MoCo动量对比机制有效工作。
消融检查：若将动量系数m设为0（即无动量更新），或队列大小K设为0，准确率应显著下降，以此验证各组件必要性。

常见问题

GPU显存不足：减小batch size至128或64，同时确保K能被batch size整除（如K=2048, batch=256）。
队列更新错误：K必须能被batch size整除，否则_dequeue_enqueue中的切片会越界。
动量编码器梯度问题：动量编码器参数必须设置requires_grad=False，否则反向传播会错误更新key编码器。

架构说明：图中展示了MoCo的核心数据流。输入图像经两种随机增强得到两个视图，分别送入查询编码器（可训练）和动量编码器（EMA更新）。查询特征与动量特征计算正样本相似度，同时与队列中4096个历史负样本特征计算相似度，共同构成InfoNCE损失。每次迭代后，当前batch的动量特征入队，最旧的batch出队，实现动态负样本库维护。