Una guía práctica para el aprendizaje contrastivo | por Mengliu Zhao | Jul, 2024

Ahora es el momento de algo aprendizaje contrastivoPara mitigar el problema de las etiquetas de anotación insuficientes y aprovechar al máximo la gran cantidad de datos sin etiquetar, se podría utilizar el aprendizaje contrastivo para ayudar de manera eficaz a la red troncal a aprender las representaciones de datos sin una tarea específica. La red troncal podría congelarse para una tarea posterior determinada y solo entrenar una red superficial en un conjunto de datos anotados limitado para lograr resultados satisfactorios.

Los enfoques de aprendizaje contrastivo más utilizados incluyen SimCLR, SimSiam y MOCO (consulte mi Artículo anterior sobre MOCO). Aquí comparamos SimCLR y SimSiam.

SimCLR calcula sobre pares positivos y negativos dentro del lote de datos, lo que requiere minería negativa estricta, pérdida de NT-Xent (que extiende la pérdida de similitud de coseno sobre un lote) y un tamaño de lote grande. SimCLR también requiere que el optimizador LARS se adapte a un tamaño de lote grande.

SimSiam, Sin embargo, utiliza una arquitectura siamesa, que evita el uso de pares negativos y, además, evita la necesidad de lotes de gran tamaño. Las diferencias entre SimSiam y SimCLR se indican en la siguiente tabla.

Comparación entre SimCLR y SimSiam. Imagen del autor.
La arquitectura de SimSiam. Fuente de la imagen: https://arxiv.org/pdf/2011.10566

En la figura anterior podemos ver que la arquitectura de SimSiam solo contiene dos partes: el codificador/columna vertebral y el predictor. Durante el tiempo de entrenamiento, se detiene la propagación del gradiente de la parte siamesa y se calcula la similitud del coseno entre las salidas de los predictores y la columna vertebral.

Entonces, ¿cómo implementamos esta arquitectura en la realidad? Continuando con el diseño de clasificación supervisada, Mantenga la columna vertebral igual y solo modifique la capa MLPEn la arquitectura de aprendizaje supervisado, el MLP genera un vector de 10 elementos que indica las probabilidades de las 10 clases. Pero para SimSiam, el propósito no es realizar una “clasificación”, sino aprender la “representación”, por lo que necesitamos que la salida sea de la misma dimensión que la salida de la red troncal para el cálculo de la pérdida. Y la similitud del coseno negativo se muestra a continuación:

import torch.nn as nn
import matplotlib.pyplot as plt

class SimSiam(nn.Module):

def __init__(self):

super(SimSiam, self).__init__()

self.backbone = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
nn.ReLU(),
nn.BatchNorm2d(32),
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
nn.ReLU(),
nn.BatchNorm2d(64),
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
nn.ReLU(),
nn.BatchNorm2d(128),
)

self.prediction_mlp = nn.Sequential(nn.Linear(128*4*4, 64),
nn.BatchNorm1d(64),
nn.ReLU(),
nn.Linear(64, 128*4*4),
)

def forward(self, x):
x = self.backbone(x)

x = x.view(-1, 128 * 4 * 4)
pred_output = self.prediction_mlp(x)
return x, pred_output

cos = nn.CosineSimilarity(dim=1, eps=1e-6)
def negative_cosine_similarity_stopgradient(pred, proj):
return -cos(pred, proj.detach()).mean()

El pseudocódigo para entrenar SimSiam se proporciona en el artículo original a continuación:

Pseudocódigo de entrenamiento para SimSiam. Fuente: https://arxiv.org/pdf/2011.10566

Y lo convertimos en código de entrenamiento real:

import tqdm

import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.transforms import RandAugment

import wandb

wandb_config = {
"learning_rate": 0.0001,
"architecture": "simsiam",
"dataset": "FashionMNIST",
"epochs": 100,
"batch_size": 256,
}

wandb.init(
# set the wandb project where this run will be logged
project="simsiam",
# track hyperparameters and run metadata
config=wandb_config,
)

# Initialize model and optimizer
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

simsiam = SimSiam()

random_augmenter = RandAugment(num_ops=5)

optimizer = optim.SGD(simsiam.parameters(),
lr=wandb_config["learning_rate"],
momentum=0.9,
weight_decay=1e-5,
)

train_dataloader = DataLoader(train_dataset, batch_size=wandb_config["batch_size"], shuffle=True)

# Training loop
for epoch in range(wandb_config["epochs"]):
simsiam.train()

print(f"Epoch {epoch}")
train_loss = 0
for batch_idx, (image, _) in enumerate(tqdm.tqdm(train_dataloader, total=len(train_dataloader))):
optimizer.zero_grad()

aug1, aug2 = random_augmenter((image*255).to(dtype=torch.uint8)).to(dtype=torch.float32) / 255.0, \
random_augmenter((image*255).to(dtype=torch.uint8)).to(dtype=torch.float32) / 255.0

proj1, pred1 = simsiam(aug1)
proj2, pred2 = simsiam(aug2)

loss = negative_cosine_similarity_stopgradient(pred1, proj2) / 2 + negative_cosine_similarity_stopgradient(pred2, proj1) / 2
loss.backward()
optimizer.step()

wandb.log({"training loss": loss})

if (epoch+1) % 10 == 0:
torch.save(simsiam.state_dict(), f"weights/simsiam_epoch{epoch+1}.pt")

Entrenamos durante 100 épocas como una comparación justa con el entrenamiento supervisado limitado; la pérdida de entrenamiento se muestra a continuación. Nota: Debido a su diseño siamés, SimSiam podría ser muy sensible a hiperparámetros como la tasa de aprendizaje y las capas ocultas de MLP. El artículo original de SimSiam proporciona una configuración detallada para la red troncal ResNet50. Para la red troncal basada en ViT, recomendamos leer el Documento MOCO v3que adopta el modelo SimSiam en un esquema de actualización de momento.

Pérdida de entrenamiento para SimSiam. Imagen del autor.

Luego, ejecutamos el SimSiam entrenado en el conjunto de prueba y visualizamos las representaciones utilizando la reducción UMAP:

import tqdm
import numpy as np

import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

simsiam = SimSiam()

test_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=False)
simsiam.load_state_dict(torch.load("weights/simsiam_epoch100.pt"))

simsiam.eval()
simsiam.to(device)

features = []
labels = []
for batch_idx, (image, target) in enumerate(tqdm.tqdm(test_dataloader, total=len(test_dataloader))):

with torch.no_grad():

proj, pred = simsiam(image.to(device))

features.extend(np.squeeze(pred.detach().cpu().numpy()).tolist())
labels.extend(target.detach().cpu().numpy().tolist())

import plotly.express as px
import umap.umap_ as umap

reducer = umap.UMAP(n_components=3, n_neighbors=10, metric="cosine")
projections = reducer.fit_transform(np.array(features))

px.scatter(projections, x=0, y=1,
color=labels, labels={'color': 'Fashion MNIST Labels'}
)

Representación UMAP de SimSiam sobre el conjunto de prueba. Imagen del autor.

Es interesante ver que hay dos pequeñas islas en el mapa de dimensión reducida que se muestra arriba: clase 5, 7, 8 y algo de 9. Si extraemos la lista de clases de FashionMNIST, sabemos que estas clases corresponden a calzado como “Sandalia”, “Zapatilla deportiva”, “Bolso” y “Botín”. El gran grupo morado corresponde a clases de ropa como “Camiseta/top”, “Pantalones”, “Suéter”, “Vestido”, “Abrigo” y “Camisa”. SimSiam demuestra el aprendizaje de una representación significativa en el dominio de la visión.