4. BYOL自举潜在表示实现

任务目标

实现BYOL（Bootstrap Your Own Latent）自监督表示学习模型，包含在线网络与目标网络，通过动量更新机制在Tiny ImageNet子集上完成预训练，并使用线性探测协议评估所学表示的质量。

核心原理与必要说明

BYOL的核心在于两个结构相同但参数不同步更新的网络：在线网络（online network）与目标网络（target network）。在线网络通过梯度下降更新，目标网络参数则通过在线网络参数的指数移动平均（EMA）获得，即动量更新：θ_target = τ·θ_target + (1−τ)·θ_online。训练时，对同一输入图像做两次不同的随机增强，分别送入两个网络，在线网络输出经预测器（predictor）后与目标网络输出的投影（projection）计算均方误差损失。由于目标网络参数滞后于在线网络，BYOL无需负样本即可避免表示坍塌。

环境依赖

需要PyTorch 1.10+、torchvision、tqdm。Tiny ImageNet数据集需预先下载并整理为ImageFolder格式（200类，每类500张训练图）。若使用完整Tiny ImageNet训练时间过长，可先使用其子集（如随机选取20类）验证代码正确性，本实现支持通过--num_classes参数控制类别数。

完整实现代码

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
from tqdm import tqdm
import copy
import argparse
import os

# ---------- 基础模块 ----------
class MLP(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        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 Encoder(nn.Module):
    """ResNet-18编码器，输出2048维特征"""
    def __init__(self):
        super().__init__()
        resnet = torchvision.models.resnet18(weights=None)
        self.features = nn.Sequential(*list(resnet.children())[:-1])
        self.fc = nn.Linear(512, 2048)
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        return self.fc(x)

class BYOL(nn.Module):
    def __init__(self, base_encoder=Encoder, proj_dim=256, pred_hidden=4096, tau=0.996):
        super().__init__()
        self.tau = tau
        # 在线网络
        self.online_encoder = base_encoder()
        self.online_projector = MLP(2048, pred_hidden, proj_dim)
        self.online_predictor = MLP(proj_dim, pred_hidden, proj_dim)
        # 目标网络（结构相同，参数不梯度更新）
        self.target_encoder = copy.deepcopy(self.online_encoder)
        self.target_projector = copy.deepcopy(self.online_projector)
        for p in self.target_encoder.parameters():
            p.requires_grad = False
        for p in self.target_projector.parameters():
            p.requires_grad = False

    @torch.no_grad()
    def _momentum_update(self):
        """动量更新目标网络参数"""
        for online_p, target_p in zip(self.online_encoder.parameters(), self.target_encoder.parameters()):
            target_p.data = self.tau * target_p.data + (1 - self.tau) * online_p.data
        for online_p, target_p in zip(self.online_projector.parameters(), self.target_projector.parameters()):
            target_p.data = self.tau * target_p.data + (1 - self.tau) * online_p.data

    def forward(self, x1, x2):
        # 在线分支
        proj1 = self.online_predictor(self.online_projector(self.online_encoder(x1)))
        # 目标分支（无梯度）
        with torch.no_grad():
            proj2 = self.target_projector(self.target_encoder(x2))
        return proj1, proj2

def byol_loss(pred, target):
    pred = F.normalize(pred, dim=1)
    target = F.normalize(target, dim=1)
    return 2 - 2 * (pred * target).sum(dim=1).mean()

# ---------- 数据增强 ----------
def get_augmentation():
    return T.Compose([
        T.RandomResizedCrop(64, scale=(0.2, 1.0)),
        T.RandomHorizontalFlip(),
        T.ColorJitter(0.4, 0.4, 0.4, 0.1),
        T.RandomGrayscale(p=0.2),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

class TwoCropsTransform:
    def __init__(self, base_transform):
        self.base = base_transform
    def __call__(self, x):
        return self.base(x), self.base(x)

# ---------- 训练 ----------
def train_byol(model, train_loader, epochs, lr, device):
    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    model.train()
    for epoch in range(epochs):
        total_loss = 0.0
        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}")
        for (x1, x2), _ in pbar:
            x1, x2 = x1.to(device), x2.to(device)
            pred, target = model(x1, x2)
            loss = byol_loss(pred, target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            model._momentum_update()
            total_loss += loss.item()
            pbar.set_postfix({"loss": f"{loss.item():.4f}"})
        scheduler.step()
        print(f"Epoch {epoch+1} avg loss: {total_loss/len(train_loader):.4f}")

# ---------- 线性评估 ----------
def linear_eval(encoder, train_loader, test_loader, device, num_classes):
    encoder.eval()
    # 提取特征
    train_feats, train_labels = [], []
    with torch.no_grad():
        for x, y in train_loader:
            x = x.to(device)
            feat = encoder(x).cpu()
            train_feats.append(feat)
            train_labels.append(y)
    train_feats = torch.cat(train_feats)
    train_labels = torch.cat(train_labels)
    # 训练线性分类器
    classifier = nn.Linear(2048, num_classes).to(device)
    opt = torch.optim.Adam(classifier.parameters(), lr=1e-3)
    ce = nn.CrossEntropyLoss()
    dataset = torch.utils.data.TensorDataset(train_feats, train_labels)
    loader = DataLoader(dataset, batch_size=256, shuffle=True)
    classifier.train()
    for _ in range(30):
        for feat, label in loader:
            feat, label = feat.to(device), label.to(device)
            out = classifier(feat)
            loss = ce(out, label)
            opt.zero_grad()
            loss.backward()
            opt.step()
    # 测试
    classifier.eval()
    correct = total = 0
    with torch.no_grad():
        for x, y in test_loader:
            x = x.to(device)
            feat = encoder(x)
            out = classifier(feat)
            pred = out.argmax(dim=1)
            correct += (pred.cpu() == y).sum().item()
            total += y.size(0)
    return correct / total

# ---------- 主程序 ----------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", type=str, required=True, help="Tiny ImageNet根目录，含train/和val/子目录")
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--batch_size", type=int, default=128)
    parser.add_argument("--lr", type=float, default=0.03)
    parser.add_argument("--num_classes", type=int, default=200, help="使用类别数，可设小值快速验证")
    parser.add_argument("--device", type=str, default="cuda")
    args = parser.parse_args()

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 数据集（需预先按类别子目录整理）
    train_transform = TwoCropsTransform(get_augmentation())
    eval_transform = T.Compose([
        T.Resize(64),
        T.CenterCrop(64),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # 若使用子集，需先筛选类别目录
    train_dir = os.path.join(args.data_dir, "train")
    val_dir = os.path.join(args.data_dir, "val")
    if args.num_classes < 200:
        classes = sorted(os.listdir(train_dir))[:args.num_classes]
        train_dir = os.path.join(args.data_dir, "train_subset")
        val_dir = os.path.join(args.data_dir, "val_subset")
        os.makedirs(train_dir, exist_ok=True)
        os.makedirs(val_dir, exist_ok=True)
        for c in classes:
            os.makedirs(os.path.join(train_dir, c), exist_ok=True)
            os.makedirs(os.path.join(val_dir, c), exist_ok=True)

    train_dataset = torchvision.datasets.ImageFolder(train_dir, transform=train_transform)
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, drop_last=True)

    # 预训练
    model = BYOL().to(device)
    train_byol(model, train_loader, args.epochs, args.lr, device)

    # 保存模型
    torch.save(model.state_dict(), "byol_tinyimagenet.pth")

    # 线性评估（使用单视图增强）
    eval_train_dataset = torchvision.datasets.ImageFolder(train_dir, transform=eval_transform)
    eval_test_dataset = torchvision.datasets.ImageFolder(val_dir, transform=eval_transform)
    eval_train_loader = DataLoader(eval_train_dataset, batch_size=256, shuffle=False, num_workers=4)
    eval_test_loader = DataLoader(eval_test_dataset, batch_size=256, shuffle=False, num_workers=4)

    acc = linear_eval(model.online_encoder, eval_train_loader, eval_test_loader, device, args.num_classes)
    print(f"Linear evaluation accuracy: {acc:.4f}")

if __name__ == "__main__":
    main()

运行方式

# 使用全部200类
python byol.py --data_dir /path/to/tiny-imagenet --epochs 50 --batch_size 128

# 快速验证（20类，约10分钟）
python byol.py --data_dir /path/to/tiny-imagenet --epochs 10 --batch_size 64 --num_classes 20

验证方法与通过标准

训练过程中观察loss应持续下降并趋于稳定（例如从初始约1.5降至0.5以下）。训练完成后，程序自动执行线性评估并输出准确率。判定通过的参考标准：在20类子集上训练10个epoch后线性评估准确率应高于50%；在完整200类上训练50个epoch后准确率应高于25%。若准确率接近随机猜测（1/类别数），说明模型发生坍塌，需检查动量更新是否生效、学习率是否过大。

常见问题

• loss不下降：检查是否调用了model._momentum_update()，以及目标网络参数是否处于requires_grad=False状态。
• 显存不足：减小batch_size或使用更小的编码器（如ResNet-18的浅层替代）。
• 数据集路径错误：Tiny ImageNet的train目录下应为200个类别子文件夹，val目录下为200个类别子文件夹（需从val_annotations.txt重新整理）。