BYOL：轻松进行自监督学习

本文作者： AI研习社-译站

2020-11-22 09:31

导语：BYOL可以利用未标记的数据来最大限度地提高模型性能。

译者：AI研习社（季一帆）

双语原文链接：Easy Self-Supervised Learning with BYOL

注：本文所有代码可见Google Colab notebook，你可用Colab的免费GPU运行或改进。

自监督学习

在深度学习中，经常遇到的问题是没有足够的标记数据，而手工标记数据耗费大量时间且人工成本高昂。基于此，自我监督学习成为深度学习的研究热点，旨在从未标记样本中进行学习，以缓解数据标注困难的问题。子监督学习的目标很简单，即训练一个模型使得相似的样本具有相似的表示，然而具体实现却困难重重。经过谷歌这样的诸多先驱者若干年的研究，子监督学习如今已取得一系列的进步与发展。

在BYOL之前，多数自我监督学习都可分为对比学习或生成学习，其中，生成学习一般GAN建模完整的数据分布，计算成本较高，相比之下，对比学习方法就很少面临这样的问题。对此，BYOL的作者这样说道：

通过对比方法，同一图像不同视图的表示更接近(正例)，不同图像视图的表示相距较远(负例)，通过这样的方式减少表示的生成成本。

为了实现对比方法，我们必须将每个样本与其他许多负例样本进行比较。然而这样会使训练很不稳定，同时会增大数据集的系统偏差。BYOL的作者显然明白这点：

对比方法对图像增强的方式非常敏感。例如，当消除图像增强中的颜色失真时，SimCLR表现不佳。可能的原因是，同一图像的不同裁切一般会共享颜色直方图，而不同图像的颜色直方图是不同的。因此，在对比任务中，可以通过关注颜色直方图，使用随机裁切方式实现图像增强，其结果表示几乎无法保留颜色直方图之外的信息。

不仅仅是颜色失真，其他类型的数据转换也是如此。一般来说，对比训练对数据的系统偏差较为敏感。在机器学习中，数据偏差是一个广泛存在的问题（见facial recognition for women and minorities），这对对比方法来说影响更大。不过好在BYOL不依赖负采样，从而很好的避免了该问题。

BYOL：Bootstrap Your Own Latent(发掘自身潜能)

BYOL的目标与对比学习相似，但一个很大的区别是，BYOL不关心不同样本是否具有不同的表征（即对比学习中的对比部分），仅仅使相似的样品表征类似。看上去似乎无关紧要，但这样的设定会显著改善模型训练效率和泛化能力：

由于不需要负采样，BLOY有更高的训练效率。在训练中，每次遍历只需对每个样本采样一次，而无需关注负样本。
BLOY模型对训练数据的系统偏差不敏感，这意味着模型可以对未见样本也有较好的适用性。

BYOL最小化样本表征和该样本变换之后的表征间的距离。其中，不同变换类型包括0：平移、旋转、模糊、颜色反转、颜色抖动、高斯噪声等（我在此以图像操作来举例说明，但BYOL也可以处理其他数据类型）。至于是单一变换还是几种不同类型的联合变换，这取决于你自己，不过我一般会采用联合变换。但有一点需要注意，如果你希望训练的模型能够应对某种变换，那么用该变换处理训练数据时必要的。

手把手教你编码BYOL

首先是数据转换增强的编码。BYOL的作者定义了一组类似于SimCLR的特殊转换：

import random
from typing import Callable, Tuple

from kornia import augmentation as aug
from kornia import filters
from kornia.geometry import transform as tf
import torch
from torch import nn, Tensor

class RandomApply(nn.Module):
def __init__(self, fn: Callable, p: float):
super().__init__()
self.fn = fn
self.p = p

def forward(self, x: Tensor) -> Tensor:
return x if random.random() > self.p else self.fn(x)

def default_augmentation(image_size: Tuple[int, int] = (224, 224)) -> nn.Module:
return nn.Sequential(
tf.Resize(size=image_size),
RandomApply(aug.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.8),
aug.RandomGrayscale(p=0.2),
aug.RandomHorizontalFlip(),
RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1),
aug.RandomResizedCrop(size=image_size),
aug.Normalize(
mean=torch.tensor([0.485, 0.456, 0.406]),
std=torch.tensor([0.229, 0.224, 0.225]),
),
)

上述代码通过Kornia实现数据转换，这是一个基于 PyTorch 的可微分的计算机视觉开源库。当然，你可以用其他开源库实现数据转换扩充，甚至是自己编写。实际上，可微分性对BYOL而言并没有那么必要。

接下来，我们编写编码器模块。该模块负责从基本模型提取特征，并将这些特征投影到低维隐空间。具体的，我们通过wrapper类实现该模块，这样我们可以轻松将BYOL用于任何模型，无需将模型编码到脚本。该类主要由两部分组成：

特征抽取，获取模型最后一层的输出。

映射，非线性层，将输出映射到更低维空间。

特征提取通过hooks实现（如果你不了解hooks，推荐阅读我之前的介绍文章How to Use PyTorch Hooks)。除此之外，代码其他部分很容易理解。

from typing import Union

def mlp(dim: int, projection_size: int = 256, hidden_size: int = 4096) -> nn.Module:
return nn.Sequential(
nn.Linear(dim, hidden_size),
nn.BatchNorm1d(hidden_size),
nn.ReLU(inplace=True),
nn.Linear(hidden_size, projection_size),
)

class EncoderWrapper(nn.Module):
def __init__(
self,
model: nn.Module,
projection_size: int = 256,
hidden_size: int = 4096,
layer: Union[str, int] = -2,
):
super().__init__()
self.model = model
self.projection_size = projection_size
self.hidden_size = hidden_size
self.layer = layer

self._projector = None
self._projector_dim = None
self._encoded = torch.empty(0)
self._register_hook()

@property
def projector(self):
if self._projector is None:
self._projector = mlp(
self._projector_dim, self.projection_size, self.hidden_size
)
return self._projector

def _hook(self, _, __, output):
output = output.flatten(start_dim=1)
if self._projector_dim is None:
self._projector_dim = output.shape[-1]
self._encoded = self.projector(output)

def _register_hook(self):
if isinstance(self.layer, str):
layer = dict([*self.model.named_modules()])[self.layer]
else:
layer = list(self.model.children())[self.layer]

layer.register_forward_hook(self._hook)

def forward(self, x: Tensor) -> Tensor:
_ = self.model(x)
return self._encoded

BYOL包含两个相同的编码器网络。第一个编码器网络的权重随着每一训练批次进行更新，而第二个网络（称为“目标”网络）使用第一个编码器权重均值进行更新。在训练过程中，目标网络接收原始批次训练数据，而另一个编码器则接收相应的转换数据。两个编码器网络会分别为相应数据生成低维表示。然后，我们使用多层感知器预测目标网络的输出，并最大化该预测与目标网络输出之间的相似性。

BYOL：轻松进行自监督学习

图源：Bootstrap Your Own Latent, Figure 2

也许有人会想，我们不是应该直接比较数据转换之前和之后的隐向量表征吗？为什么还有设计多层感知机？假设没有MLP层的话，网络可以通过将权重降低到零方便的使所有图像的表示相似化，可这样模型并没有学到任何有用的东西，而MLP层可以识别出数据转换并预测目标隐向量。这样避免了权重趋零，可以学习更恰当的数据表示！

训练结束后，舍弃目标网络编码器，只保留一个编码器，根据该编码器，所有训练数据可生成自洽表示。这正是BYOL能够进行自监督学习的关键！因为学习到的表示具有自洽性，所以经不同的数据变换后几乎保持不变。这样，模型使得相似示例的表示更加接近！

接下来编写BYOL的训练代码。我选择使用Pythorch Lightning开源库，该库基于PyTorch，对深度学习项目非常友好，能够进行多GPU培训、实验日志记录、模型断点检查和混合精度训练等，甚至在cloud TPU上也支持基于该库运行PyTorch模型！

from copy import deepcopy
from itertools import chain
from typing import Dict, List

import pytorch_lightning as pl
from torch import optim
import torch.nn.functional as f

def normalized_mse(x: Tensor, y: Tensor) -> Tensor:
x = f.normalize(x, dim=-1)
y = f.normalize(y, dim=-1)
return 2 - 2 * (x * y).sum(dim=-1)

class BYOL(pl.LightningModule):
def __init__(
self,
model: nn.Module,
image_size: Tuple[int, int] = (128, 128),
hidden_layer: Union[str, int] = -2,
projection_size: int = 256,
hidden_size: int = 4096,
augment_fn: Callable = None,
beta: float = 0.99,
**hparams,
):
super().__init__()
self.augment = default_augmentation(image_size) if augment_fn is None else augment_fn
self.beta = beta
self.encoder = EncoderWrapper(
model, projection_size, hidden_size, layer=hidden_layer
)
self.predictor = nn.Linear(projection_size, projection_size, hidden_size)
self.hparams = hparams
self._target = None

self.encoder(torch.zeros(2, 3, *image_size))

def forward(self, x: Tensor) -> Tensor:
return self.predictor(self.encoder(x))

@property
def target(self):
if self._target is None:
self._target = deepcopy(self.encoder)
return self._target

def update_target(self):
for p, pt in zip(self.encoder.parameters(), self.target.parameters()):
pt.data = self.beta * pt.data + (1 - self.beta) * p.data

# --- Methods required for PyTorch Lightning only! ---

def configure_optimizers(self):
optimizer = getattr(optim, self.hparams.get("optimizer", "Adam"))
lr = self.hparams.get("lr", 1e-4)
weight_decay = self.hparams.get("weight_decay", 1e-6)
return optimizer(self.parameters(), lr=lr, weight_decay=weight_decay)

def training_step(self, batch, *_) -> Dict[str, Union[Tensor, Dict]]:
x = batch[0]
with torch.no_grad():
x1, x2 = self.augment(x), self.augment(x)

pred1, pred2 = self.forward(x1), self.forward(x2)
with torch.no_grad():
targ1, targ2 = self.target(x1), self.target(x2)
loss = torch.mean(normalized_mse(pred1, targ2) + normalized_mse(pred2, targ1))

self.log("train_loss", loss.item())
return {"loss": loss}

@torch.no_grad()
def validation_step(self, batch, *_) -> Dict[str, Union[Tensor, Dict]]:
x = batch[0]
x1, x2 = self.augment(x), self.augment(x)
pred1, pred2 = self.forward(x1), self.forward(x2)
targ1, targ2 = self.target(x1), self.target(x2)
loss = torch.mean(normalized_mse(pred1, targ2) + normalized_mse(pred2, targ1))

return {"loss": loss}

@torch.no_grad()
def validation_epoch_end(self, outputs: List[Dict]) -> Dict:
val_loss = sum(x["loss"] for x in outputs) / len(outputs)
self.log("val_loss", val_loss.item())

上述代码部分源自Pythorch Lightning提供的示例代码。这段代码你尤其需要关注的是training_step，在此函数实现模型的数据转换、特征投影和相似性损失计算等。

实例说明

下文我们将在STL10数据集上对BYOL进行实验验证。因为该数据集同时包含大量未标记的图像以及标记的训练和测试集，非常适合无监督和自监督学习实验。STL10网站这样描述该数据集：

STL-10数据集是一个用于研究无监督特征学习、深度学习、自学习算法的图像识别数据集。该数据集是对CIFAR-10数据集的改进，最明显的便是，每个类的标记训练数据比CIFAR-10中的要少，但在监督训练之前，数据集提供大量的未标记样本训练模型学习图像模型。因此，该数据集主要的挑战是利用未标记的数据（与标记数据相似但分布不同）来构建有用的先验知识。

通过Torchvision可以很方便的加载STL10，因此无需担心数据的下载和预处理。

from torchvision.datasets import STL10
from torchvision.transforms import ToTensor

TRAIN_DATASET = STL10(root="data", split="train", download=True, transform=ToTensor())
TRAIN_UNLABELED_DATASET = STL10(
root="data", split="train+unlabeled", download=True, transform=ToTensor()
)
TEST_DATASET = STL10(root="data", split="test", download=True, transform=ToTensor())

同时，我们使用监督学习方法作为基准模型，以此衡量本文模型的准确性。基线模型也可通过Lightning模块轻易实现：

class SupervisedLightningModule(pl.LightningModule):
def __init__(self, model: nn.Module, **hparams):
super().__init__()
self.model = model

def forward(self, x: Tensor) -> Tensor:
return self.model(x)

def configure_optimizers(self):
optimizer = getattr(optim, self.hparams.get("optimizer", "Adam"))
lr = self.hparams.get("lr", 1e-4)
weight_decay = self.hparams.get("weight_decay", 1e-6)
return optimizer(self.parameters(), lr=lr, weight_decay=weight_decay)

def training_step(self, batch, *_) -> Dict[str, Union[Tensor, Dict]]:
x, y = batch
loss = f.cross_entropy(self.forward(x), y)
self.log("train_loss", loss.item())
return {"loss": loss}

@torch.no_grad()
def validation_step(self, batch, *_) -> Dict[str, Union[Tensor, Dict]]:
x, y = batch
loss = f.cross_entropy(self.forward(x), y)
return {"loss": loss}

@torch.no_grad()
def validation_epoch_end(self, outputs: List[Dict]) -> Dict:
val_loss = sum(x["loss"] for x in outputs) / len(outputs)
self.log("val_loss", val_loss.item())

可以看到，使用Pythorch Lightning可以方便的构建并训练模型。只需为训练集和测试集创建DataLoader对象，将其导入需要训练的模型即可。本实验中，epoch设置为25，学习率为1e-4。

from os import cpu_count

from torch.utils.data import DataLoader
from torchvision.models import resnet18

model = resnet18(pretrained=True)
supervised = SupervisedLightningModule(model)
trainer = pl.Trainer(max_epochs=25, gpus=-1, weights_summary=None)
train_loader = DataLoader(
TRAIN_DATASET,
batch_size=128,
shuffle=True,
drop_last=True,
)
val_loader = DataLoader(
TEST_DATASET,
batch_size=128,
)
trainer.fit(supervised, train_loader, val_loader)

经训练，仅通过一个非常小的模型ResNet18就取得约85%的准确率。但实际上，我们还可以做得更好！

接下来，我们使用BYOL对ResNet18模型进行预训练。在这次实验中，我选择epoch为50，学习率依然是1e-4。注：该过程是本文代码耗时最长的部分，在K80 GPU的标准Colab中大约需要45分钟。

model = resnet18(pretrained=True)
byol = BYOL(model, image_size=(96, 96))
trainer = pl.Trainer(
max_epochs=50,
gpus=-1,
accumulate_grad_batches=2048 // 128,
weights_summary=None,
)
train_loader = DataLoader(
TRAIN_UNLABELED_DATASET,
batch_size=128,
shuffle=True,
drop_last=True,
)
trainer.fit(byol, train_loader, val_loader)

然后，我们使用新的ResNet18模型重新进行监督学习。（为彻底清除BYOL中的前向hook，我们实例化一个新模型，在该模型引入经过训练的状态字典。）

# Extract the state dictionary, initialize a new ResNet18 model,
# and load the state dictionary into the new model.
#
# This ensures that we remove all hooks from the previous model,
# which are automatically implemented by BYOL.
state_dict = model.state_dict()
model = resnet18()
model.load_state_dict(state_dict)

supervised = SupervisedLightningModule(model)
trainer = pl.Trainer(
max_epochs=25,
gpus=-1,
weights_summary=None,
)
train_loader = DataLoader(
TRAIN_DATASET,
batch_size=128,
shuffle=True,
drop_last=True,
)
trainer.fit(supervised, train_loader, val_loader)

通过这种方式，模型准确率提高了约2.5%，达到了87.7%！虽然该方法需要更多的代码（大约300行）以及一些库的支撑，但相比其他自监督方法仍显得简洁。作为对比，可以看下官方的SimCLR或SwAV是多么复杂。而且，本文具有更快的训练速度，即使是Colab的免费GPU，整个实验也不到一个小时。

结论

本文要点总结如下。首先也是最重要的，BYOL是一种巧妙的自监督学习方法，可以利用未标记的数据来最大限度地提高模型性能。此外，由于所有ResNet模型都是使用ImageNet进行预训练的，因此BYOL的性能优于预训练的ResNet18。STL10是ImageNet的一个子集，所有图像都从224x224像素缩小到96x96像素。虽然分辨率发生改变，我们希望自监督学习能避免这样的影响，表现出较好性能，而仅仅依靠STL10的小规模训练集是不够的。

类似ResNet这样的模型中，ML从业人员过于依赖预先训练的权重。虽然这在一定情况下是很好的选择，但不一定适合其他数据，哪怕在STL10这样与ImageNet高度相似的数据中表现也不如人意。因此，我迫切希望将来在深度学习的研究中，自监督方法能够获得更多的关注与实践应用。