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

1次阅读

没有评论

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

自监督学习是机器学习领域的重要分支，它通过从无标签数据中自动生成监督信号来训练模型。对比学习作为自监督学习的一种方法，其核心思想是通过最大化相似样本之间的相似度，同时最小化不相似样本之间的相似度来学习特征表示。

传统对比学习方法如 SimCLR 依赖于大量的负样本来防止模型坍塌（即所有样本映射到同一个点），但这种方式在计算和存储上成本较高。BYOL（Bootstrap Your Own Latent）则创新性地提出了一种无需负样本的对比学习框架。

BYOL 的主要创新在于它通过两个神经网络（在线网络和目标网络）的交互来避免使用负样本。具体来说：

在线网络 ：通过梯度下降更新参数，负责学习数据表示。
目标网络 ：通过指数移动平均（EMA）更新参数，提供稳定的目标表示。

BYOL 通过最小化在线网络和目标网络对同一数据的不同增强视图的预测误差来实现特征学习，从而避免了负样本的需求。

BYOL 包含两个结构相同的网络：

在线网络 ：由编码器 $f_\theta$、投影头 $g_\theta$ 和预测头 $q_\theta$ 组成，参数通过梯度下降更新。
目标网络 ：由编码器 $f_\xi$ 和投影头 $g_\xi$ 组成，参数通过 EMA 更新：$\xi \leftarrow \tau \xi + (1-\tau)\theta$，其中 $\tau$ 是动量系数。

预测头 $q_\theta$ 是一个小型 MLP，用于将在线网络的输出映射到目标网络的空间，增强表示的学习能力。

EMA 更新确保了目标网络的参数变化平滑，避免了在线网络的快速变化导致训练不稳定。

以下是 BYOL 的完整 PyTorch 实现代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms

# 数据增强
class Augmentation:
    def __init__(self):
        self.transform = transforms.Compose([transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(0.4, 0.4, 0.4, 0.1),
            transforms.GaussianBlur(kernel_size=23, sigma=(0.1, 2.0)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, x):
        return self.transform(x), self.transform(x)

# 编码器（以 ResNet 为例）class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        # 其余层省略...

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        return x

# 投影头
class ProjectionHead(nn.Module):
    def __init__(self, input_dim=2048, hidden_dim=512, output_dim=256):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.bn1 = nn.BatchNorm1d(hidden_dim)
        self.relu = nn.ReLU(inplace=True)
        self.layer2 = nn.Linear(hidden_dim, output_dim)

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

# 预测头
class PredictionHead(nn.Module):
    def __init__(self, input_dim=256, hidden_dim=512, output_dim=256):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.bn1 = nn.BatchNorm1d(hidden_dim)
        self.relu = nn.ReLU(inplace=True)
        self.layer2 = nn.Linear(hidden_dim, output_dim)

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

# BYOL 模型
class BYOL(nn.Module):
    def __init__(self, encoder, hidden_dim=512, projection_dim=256, tau=0.996):
        super().__init__()
        self.tau = tau

        # 在线网络
        self.online_encoder = encoder
        self.online_projector = ProjectionHead(output_dim=projection_dim)
        self.online_predictor = PredictionHead(output_dim=projection_dim)

        # 目标网络
        self.target_encoder = encoder
        self.target_projector = ProjectionHead(output_dim=projection_dim)

        # 初始化目标网络参数与在线网络相同
        self._init_target_network()

    def _init_target_network(self):
        for online_param, target_param in zip(self.online_encoder.parameters(), self.target_encoder.parameters()):
            target_param.data.copy_(online_param.data)
            target_param.requires_grad = False

        for online_param, target_param in zip(self.online_projector.parameters(), self.target_projector.parameters()):
            target_param.data.copy_(online_param.data)
            target_param.requires_grad = False

    def update_target_network(self):
        # EMA 更新目标网络
        for online_param, target_param in zip(self.online_encoder.parameters(), self.target_encoder.parameters()):
            target_param.data = self.tau * target_param.data + (1 - self.tau) * online_param.data

        for online_param, target_param in zip(self.online_projector.parameters(), self.target_projector.parameters()):
            target_param.data = self.tau * target_param.data + (1 - self.tau) * online_param.data

    def forward(self, x1, x2):
        # 在线网络处理第一个视图
        online_z1 = self.online_projector(self.online_encoder(x1))
        online_q1 = self.online_predictor(online_z1)

        # 在线网络处理第二个视图
        online_z2 = self.online_projector(self.online_encoder(x2))
        online_q2 = self.online_predictor(online_z2)

        # 目标网络处理第一个视图
        with torch.no_grad():
            target_z1 = self.target_projector(self.target_encoder(x1))

        # 目标网络处理第二个视图
        with torch.no_grad():
            target_z2 = self.target_projector(self.target_encoder(x2))

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

# 对称损失函数
def byol_loss(q, z):
    q = F.normalize(q, dim=-1)
    z = F.normalize(z, dim=-1)
    return 2 - 2 * (q * z).sum(dim=-1).mean()