BYOL对比学习原理解析与自监督实践指南

1次阅读

没有评论

共计 4056 个字符，预计需要花费 11 分钟才能阅读完成。

自监督学习近年来成为深度学习领域的热门方向，它通过设计巧妙的预训练任务，从无标注数据中学习有用的特征表示。对比学习作为自监督学习的重要分支，其核心思想是通过拉近相似样本（正样本对）的特征表示，同时推远不相似样本（负样本对）的特征表示。

然而，传统的对比学习方法如 SimCLR 面临两个主要挑战：

负样本依赖 ：需要大量负样本才能获得良好的特征表示，这导致计算开销巨大
样本冲突 ：随机采样的负样本可能包含与正样本语义相似的样本，影响学习效果

BYOL(Bootstrap Your Own Latent) 通过创新的网络架构和训练机制，成功摆脱了对负样本的依赖，成为自监督学习领域的重要突破。

BYOL 包含两个结构相同但参数不同的网络：

在线网络 (online network)：包含 encoder $f_θ$、projector $g_θ$ 和 predictor $q_θ$，通过梯度下降更新参数
目标网络 (target network)：包含 encoder $f_ξ$ 和 projector $g_ξ$，其参数通过动量更新从在线网络获得

这种设计使得目标网络能够提供稳定的监督信号，避免了传统对比学习中需要大量负样本的问题。

目标网络的参数更新遵循动量更新规则：

$$ξ ← τξ + (1-τ)θ$$

其中 $τ$ 是动量系数，通常设置为 0.99。这种平滑更新方式保证了目标网络的变化不会过于剧烈，为在线网络提供了稳定的学习目标。

预测头 $q_θ$ 是 BYOL 的关键设计之一，它：

增加了网络的表达能力
防止网络陷入简单的解 (如将所有样本映射到同一点)
帮助在线网络学习更丰富的特征表示

数学上，BYOL 的损失函数定义为：

$$L_{θ,ξ} = ||q_θ(z_θ) – z’_ξ||_2^2$$

其中 $z_θ = g_θ(f_θ(x))$, $z’_ξ = g_ξ(f_ξ(x’))$，$x$ 和 $x’$ 是同一图像的两个不同增强视图。

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

class MLPHead(nn.Module):
    """投影头和预测头"""
    def __init__(self, in_dim, hidden_dim=512, out_dim=128):
        super().__init__()
        self.layer1 = nn.Sequential(nn.Linear(in_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True)
        )
        self.layer2 = nn.Linear(hidden_dim, out_dim)

    def forward(self, x):
        x = self.layer1(x)
        return self.layer2(x)

class BYOL(nn.Module):
    """BYOL 主网络"""
    def __init__(self, backbone, hidden_dim=512, out_dim=128, momentum=0.99):
        super().__init__()
        self.momentum = momentum

        # 在线网络
        self.online_encoder = backbone
        self.online_projector = MLPHead(backbone.output_dim, hidden_dim, out_dim)
        self.online_predictor = MLPHead(out_dim, hidden_dim, out_dim)

        # 目标网络 (初始时与在线网络相同)
        self.target_encoder = copy.deepcopy(backbone)
        self.target_projector = copy.deepcopy(self.online_projector)

        # 冻结目标网络参数
        for param in self.target_encoder.parameters():
            param.requires_grad = False
        for param in self.target_projector.parameters():
            param.requires_grad = False

    @torch.no_grad()
    def update_target(self):
        """动量更新目标网络"""
        for online, target in zip([self.online_encoder, self.online_projector],
            [self.target_encoder, self.target_projector]
        ):
            for o_param, t_param in zip(online.parameters(), target.parameters()):
                t_param.data = t_param.data * self.momentum + o_param.data * (1. - self.momentum)

    def forward(self, x1, x2):
        """
        输入:
            x1, x2: 同一图像的两个不同增强视图
        返回:
            online_view1: 在线网络对 x1 的输出
            online_view2: 在线网络对 x2 的输出
            target_view1: 目标网络对 x1 的输出
            target_view2: 目标网络对 x2 的输出
        """
        # 在线网络处理两个视图
        online_z1 = self.online_projector(self.online_encoder(x1))
        online_z2 = self.online_projector(self.online_encoder(x2))
        online_q1 = self.online_predictor(online_z1)
        online_q2 = self.online_predictor(online_z2)

        # 目标网络处理两个视图 (不计算梯度)
        with torch.no_grad():
            target_z1 = self.target_projector(self.target_encoder(x1))
            target_z2 = self.target_projector(self.target_encoder(x2))

        return online_q1, online_q2, target_z1.detach(), target_z2.detach()

BYOL 的性能高度依赖数据增强，以下是关键增强操作：

from torchvision import transforms

train_transform = transforms.Compose([transforms.RandomResizedCrop(size=32, scale=(0.08, 1.0)),  # 随机裁剪并缩放
    transforms.RandomHorizontalFlip(p=0.5),  # 随机水平翻转
    transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1),  # 颜色抖动
    transforms.RandomGrayscale(p=0.2),  # 随机灰度化
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010])
])

def byol_loss(online, target):
    """计算 BYOL 对称损失"""
    # 归一化处理
    online = F.normalize(online, dim=-1)
    target = F.normalize(target, dim=-1)
    # 计算均方误差
    return 2 - 2 * (online * target).sum(dim=-1)

# 使用示例
online_q1, online_q2, target_z1, target_z2 = model(x1, x2)
loss1 = byol_loss(online_q1, target_z2)
loss2 = byol_loss(online_q2, target_z1)
loss = (loss1 + loss2).mean()

硬件配置 ：NVIDIA V100 GPU, 32GB 显存
数据集 ：CIFAR-10 训练集 (50,000 张图像)
模型：ResNet-18 作为 backbone
训练参数 ：
batch size: 512
学习率: 0.06(带 cosine 衰减)
训练 epoch: 200
动量系数: 0.99