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高度相似的數據中表現也不如人意。因此，我迫切希望將來在深度學習的研究中，自監督方法能夠獲得更多的關注與實踐應用。