BYOL对比学习入门指南：从理论到PyTorch实战

1次阅读

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

自监督学习是近年来机器学习领域的重要方向，它通过从无标注数据中自动生成监督信号来训练模型。对比学习作为自监督学习的核心范式，其核心思想是拉近相似样本（正样本对）的表示距离，同时推开不相似样本（负样本对）的表示距离。

传统对比学习方法（如 SimCLR）依赖于大量负样本构建对比任务，但负样本的采集和存储会带来计算开销和内存压力。BYOL（Bootstrap Your Own Latent）的创新之处在于完全摒弃负样本，仅通过在线网络和目标网络的交互实现特征学习。

BYOL 与 SimCLR 等传统方法的区别主要体现在三个方面：

无需负样本 ：通过目标网络提供的稳定表征作为学习目标，避免负样本带来的计算负担
非对称结构 ：在线网络包含预测头，目标网络通过 EMA（指数移动平均）更新
目标一致性 ：目标网络参数缓慢更新，为在线网络提供一致的学习目标

其核心公式可表示为：
$$\mathcal{L}{\theta,\xi} = |q\theta(z_\theta) – z’\xi|_2^2$$
其中 $q\theta$ 是在线网络的预测头，$z_\theta$ 和 $z’_\xi$ 分别是在线网络和目标网络的输出。

在线网络 ：包含编码器 $f_\theta$ 和预测头 $q_\theta$，通过梯度下降更新
目标网络 ：结构同在线网络（不含预测头），参数 $\xi$ 通过 EMA 更新：
$$\xi \leftarrow \tau\xi + (1-\tau)\theta$$
典型 $\tau$ 值取 0.99-0.999

BYOL 对同一输入应用两种随机增强：

全局视图：包含颜色失真、高斯模糊等强增强
局部视图：随机裁剪等空间变换

对 batch 内每个样本计算两个方向的预测误差：
$$\mathcal{L} = \frac{1}{2N}\sum_{i=1}^N (|q(z_i)-z’_i|^2 + |q(z’_i)-z_i|^2)$$

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

class MLPHead(nn.Module):
    """BYOL 的预测头"""
    def __init__(self, dim_in, hidden_size=4096, dim_out=256):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(dim_in, hidden_size),
            nn.BatchNorm1d(hidden_size),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_size, dim_out)
        )

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

class BYOL(nn.Module):
    def __init__(self, backbone, hidden_size=4096, pred_dim=256, tau=0.996):
        super().__init__()
        self.tau = tau
        self.online_encoder = backbone
        self.target_encoder = backbone

        # 初始化目标网络与在线网络参数一致
        for param_o, param_t in zip(self.online_encoder.parameters(), 
                                   self.target_encoder.parameters()):
            param_t.data.copy_(param_o.data)
            param_t.requires_grad = False  # 目标网络不更新梯度

        # 预测头
        self.online_predictor = MLPHead(dim_in=backbone.output_dim, 
                                      hidden_size=hidden_size,
                                      dim_out=pred_dim)

    @torch.no_grad()
    def update_target(self):
        """EMA 更新目标网络"""
        for param_o, param_t in zip(self.online_encoder.parameters(),
                                   self.target_encoder.parameters()):
            param_t.data = self.tau * param_t.data + (1 - self.tau) * param_o.data

    def forward(self, view1, view2):
        """
        输入: 
            view1, view2 - 同一图像的两个增强视图
        返回:
            loss - 对称损失值
        """
        # 在线网络前向
        z1_o = self.online_encoder(view1)
        z2_o = self.online_encoder(view2)
        p1_o = self.online_predictor(z1_o)
        p2_o = self.online_predictor(z2_o)

        # 目标网络前向 (不计算梯度)
        with torch.no_grad():
            z2_t = self.target_encoder(view2)
            z1_t = self.target_encoder(view1)

        # 对称损失
        loss = 0.5 * (F.mse_loss(p1_o, z2_t.detach()) + 
                     F.mse_loss(p2_o, z1_t.detach()))

        # 更新目标网络
        self.update_target()
        return loss