Augmentación y self-supervision

El problema

Tienes 40 imágenes en escala de grises de 11×11 píxeles. Cada una contiene una flecha horizontal centrada con un poco de ruido. La etiqueta dice hacia dónde apunta: 0 para izquierda y 1 para derecha. Hay 20 de cada clase.

Un MLP entrenado con esas 40 imágenes acierta todas. En un conjunto de validación de 200 imágenes, las flechas aparecen desplazadas hasta 2 píxeles en cada dirección, y el mismo modelo acierta solo 0.585.

Conseguir más imágenes etiquetadas no es posible. Un compañero propone generar más ejemplos a partir de los 40 que ya tienes: voltearlos, moverlos o añadirles ruido, y conservar la etiqueta original.

Transformaciones de tensor

La augmentación de datos crea ejemplos de entrenamiento nuevos aplicando transformaciones a los que ya existen. Con una imagen guardada como tensor, tres transformaciones básicas son operaciones de PyTorch:

• torch.flip(x, dims=[d]) invierte el orden de los elementos en la dimensión d
• torch.roll(x, shifts=k, dims=d) desplaza los elementos k posiciones en la dimensión d
• x + sigma * torch.randn_like(x) suma ruido gaussiano con desviación estándar sigma

En una imagen 2D, la dimensión 0 recorre filas y la dimensión 1 recorre columnas. La celda plegada importa PyTorch y define flecha, una imagen de 7×7 con unos en los trazos y ceros en el fondo.

import math

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

flecha = torch.tensor(
    [
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0, 1, 0],
        [0, 1, 1, 1, 1, 1, 1],
        [0, 0, 0, 0, 0, 1, 0],
        [0, 0, 0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
    ],
    dtype=torch.float32,
)

.int() convierte los valores a enteros solo para que la impresión sea compacta.

print(flecha.int())

tensor([[0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0, 1, 0],
        [0, 1, 1, 1, 1, 1, 1],
        [0, 0, 0, 0, 0, 1, 0],
        [0, 0, 0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0, 0, 0]], dtype=torch.int32)

flecha apunta a la derecha: el trazo horizontal ocupa la fila 3 y la punta termina en la última columna. torch.equal devuelve True si dos tensores tienen la misma forma y los mismos valores.

Pregunta. Con esta flecha, dibuja el resultado de voltear en la dimensión 1. ¿El volteo en la dimensión 0 produce una imagen distinta? ¿En qué fila queda el trazo horizontal después de torch.roll(flecha, shifts=1, dims=0)?

volteo_columnas = torch.flip(flecha, dims=[1])
volteo_filas = torch.flip(flecha, dims=[0])
bajada = torch.roll(flecha, shifts=1, dims=0)

print(volteo_columnas.int())
print(torch.equal(volteo_filas, flecha))
print(bajada.int())

tensor([[0, 0, 0, 0, 0, 0, 0],
        [0, 0, 1, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 0],
        [0, 1, 0, 0, 0, 0, 0],
        [0, 0, 1, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0]], dtype=torch.int32)
True
tensor([[0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0, 1, 0],
        [0, 1, 1, 1, 1, 1, 1],
        [0, 0, 0, 0, 0, 1, 0],
        [0, 0, 0, 0, 1, 0, 0]], dtype=torch.int32)

torch.roll es circular: lo que sale por un borde entra por el opuesto. Aquí la última fila era de ceros, así que la fila que reaparece arriba no añade trazos. Con un desplazamiento mayor que el margen vacío, parte de la flecha aparecería del otro lado.

Transformación y etiqueta

Una transformación sirve como augmentación solo si la etiqueta sigue siendo verdadera para la imagen transformada. Eso depende de la tarea, no de la imagen.

Las imágenes del problema se construyen colocando la flecha en un lienzo de 11×11. lienzo(1) pone la flecha que apunta a la derecha y lienzo(0), la que apunta a la izquierda. Los argumentos dy y dx la mueven. La celda plegada define lienzo y guarda las dos flechas centradas en derecha e izquierda.

def lienzo(direccion, dy=0, dx=0):
    imagen = torch.zeros(11, 11)
    base = flecha if direccion == 1 else torch.flip(flecha, dims=[1])
    # margen de 2 píxeles en cada borde cuando dy = dx = 0
    imagen[2 + dy : 9 + dy, 2 + dx : 9 + dx] = base
    return imagen


derecha = lienzo(1)
izquierda = lienzo(0)

En derecha, la flecha de 7×7 ocupa las filas y columnas 2 a 8 del lienzo. Sus trazos ocupan las filas 3 a 7 y las columnas 3 a 8, así que a la derecha de la punta quedan 2 columnas vacías. izquierda contiene la flecha que apunta a la izquierda en la misma posición.

Además del volteo y el desplazamiento, la celda siguiente prueba torch.rot90, que gira una imagen 90° en sentido antihorario. Con tuplas en shifts y dims, torch.roll desplaza filas y columnas en una sola llamada.

Pregunta. Ya viste qué imagen produce cada volteo; ahora decide qué pasa con la etiqueta. Para cada transformación de la flecha derecha, decide si la etiqueta 1 sigue siendo correcta, si la imagen pasa a ser de la clase 0 o si ya no pertenece a ninguna de las dos clases.

print(torch.equal(torch.flip(derecha, dims=[0]), derecha))
print(torch.equal(torch.flip(derecha, dims=[1]), izquierda))
print(torch.equal(torch.roll(derecha, shifts=(1, 2), dims=(0, 1)), lienzo(1, 1, 2)))
print(torch.rot90(flecha).int())

True
True
True
tensor([[0, 0, 0, 1, 0, 0, 0],
        [0, 0, 1, 1, 1, 0, 0],
        [0, 1, 0, 1, 0, 1, 0],
        [0, 0, 0, 1, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0]], dtype=torch.int32)

El ruido tiene un límite distinto: no convierte una clase en otra, pero puede borrar la información que define la etiqueta. Para medirlo sin un modelo entrenado, se suman 200 lienzos de ruido distintos a derecha con cada sigma, se compara cada imagen ruidosa con las dos flechas limpias y se elige la más cercana en distancia euclidiana.

Pregunta. Para 200 flechas derechas con ruido, ¿qué fracción queda más cerca de derecha con sigma = 0.1? ¿Y con sigma = 2.0?

Augmentación en entrenamiento

Ahora se construyen los datos del problema. conjunto genera imágenes con clases alternadas, un desplazamiento aleatorio de hasta max_desp píxeles y ruido con sigma = 0.1. Las imágenes tienen forma (N, 1, 11, 11): ejemplo, canal, filas y columnas. La celda plegada define conjunto y genera con ella X_ent, 40 imágenes centradas (max_desp=0), y X_val, 200 imágenes desplazadas hasta 2 píxeles (max_desp=2).

def conjunto(n, max_desp, semilla):
    gen = torch.Generator().manual_seed(semilla)
    imagenes, etiquetas = [], []
    for i in range(n):
        clase = i % 2
        dy, dx = torch.randint(-max_desp, max_desp + 1, (2,), generator=gen).tolist()
        imagen = lienzo(clase, dy, dx) + 0.1 * torch.randn(11, 11, generator=gen)
        imagenes.append(imagen)
        etiquetas.append(clase)
    return torch.stack(imagenes).unsqueeze(1), torch.tensor(etiquetas)


X_ent, y_ent = conjunto(40, max_desp=0, semilla=1)
X_val, y_val = conjunto(200, max_desp=2, semilla=2)

X_ent tiene forma (40, 1, 11, 11) y X_val tiene forma (200, 1, 11, 11). y_ent tiene 20 imágenes de la clase 0 y 20 de la clase 1.

La augmentación se aplica dentro del loop de entrenamiento. En cada época, el batch pasa por transformaciones aleatorias nuevas, así que el modelo ve una variante distinta de cada imagen. desplazar mueve cada imagen hasta 2 píxeles en filas y columnas con torch.roll. Un torch.Generator con semilla propia, el argumento gen, fija los números aleatorios de las funciones que lo reciben con generator=. torch.empty_like(x) reserva un tensor de la misma forma que x para guardar el resultado.

def desplazar(x, gen, maximo=2):
    salida = torch.empty_like(x)
    for i in range(len(x)):
        dy, dx = torch.randint(-maximo, maximo + 1, (2,), generator=gen).tolist()
        # x[i] tiene forma (C, H, W): filas y columnas son las dimensiones 1 y 2
        salida[i] = torch.roll(x[i], shifts=(dy, dx), dims=(1, 2))
    return salida

voltear aplica el volteo de columnas a cada imagen con probabilidad 0.5; en un tensor (N, C, H, W), las columnas son la dimensión 3. torch.rand(len(x)) < p sortea una máscara booleana con un valor por imagen, y x.clone() copia el batch para no modificar X_ent.

def voltear(x, gen, p=0.5):
    elegidas = torch.rand(len(x), generator=gen) < p
    salida = x.clone()
    salida[elegidas] = torch.flip(x[elegidas], dims=[3])
    return salida

entrenar ajusta un MLP de 121 entradas, 32 unidades ocultas y 2 salidas durante 150 épocas, con las 40 imágenes en cada paso. El argumento aumentar recibe una función de augmentación; si se pasa, se aplica a X_ent en cada época antes de calcular la pérdida. Ninguna de las dos funciones toca y_ent.

def entrenar(aumentar=None, epocas=150):
    torch.manual_seed(0)  # mismos pesos iniciales en todas las configuraciones
    modelo = nn.Sequential(
        nn.Flatten(), nn.Linear(121, 32), nn.ReLU(), nn.Linear(32, 2)
    )
    optimizador = torch.optim.Adam(modelo.parameters(), lr=0.01)
    gen = torch.Generator().manual_seed(0)
    for _ in range(epocas):
        modelo.train()
        x = X_ent if aumentar is None else aumentar(X_ent, gen)
        perdida = F.cross_entropy(modelo(x), y_ent)
        optimizador.zero_grad()
        perdida.backward()
        optimizador.step()
    return modelo

La evaluación usa X_val tal como está. exactitud, en la celda plegada, pone el modelo en modo de evaluación y devuelve la fracción de aciertos con tres decimales. La tercera configuración compone las dos augmentaciones con una lambda: primero desplaza y después voltea.

def exactitud(modelo, x, y):
    modelo.eval()
    with torch.no_grad():
        return round((modelo(x).argmax(dim=1) == y).float().mean().item(), 3)

Pregunta. Ordena las tres configuraciones de la celda según su exactitud de validación: sin augmentación, con desplazamiento y con desplazamiento más volteo.

modelos = {
    "sin augmentación": entrenar(),
    "desplazamiento": entrenar(desplazar),
    "desplazamiento + volteo": entrenar(lambda x, gen: voltear(desplazar(x, gen), gen)),
}

for nombre, modelo in modelos.items():
    print(nombre, exactitud(modelo, X_ent, y_ent), exactitud(modelo, X_val, y_val))

sin augmentación 1.0 0.585
desplazamiento 1.0 1.0
desplazamiento + volteo 0.45 0.4

Pregunta. Evalúas dos veces el modelo con desplazamiento sobre X_val transformado con desplazar. ¿Las dos exactitudes coinciden? ¿Coinciden con la exactitud sobre X_val sin transformar?

modelo_desp = modelos["desplazamiento"]
gen_eval = torch.Generator().manual_seed(5)

print(exactitud(modelo_desp, desplazar(X_val, gen_eval), y_val))
print(exactitud(modelo_desp, desplazar(X_val, gen_eval), y_val))
print(exactitud(modelo_desp, X_val, y_val))

0.835
0.83
1.0

Dos vistas

En clasificación, la augmentación conserva una etiqueta que ya existe. La self-supervision obtiene la señal de entrenamiento de los propios datos, sin etiquetas. Una de sus formas, el aprendizaje contrastivo, usa las transformaciones. De una imagen se obtienen dos vistas con transformaciones aleatorias. Un encoder f convierte cada imagen en un vector, su representación. El entrenamiento pide que las dos vistas de la misma imagen tengan representaciones cercanas y que las vistas de imágenes distintas queden lejos. Nadie dice qué contiene cada imagen; la señal de entrenamiento sale de saber qué vistas vienen de la misma imagen.

La cercanía se mide con la similitud coseno: cos(u, v) = u · v / (‖u‖ ‖v‖). Vale 1 si los vectores apuntan en la misma dirección, 0 si son perpendiculares y −1 si apuntan en direcciones opuestas.

Antes de entrenar un encoder, la representación más simple es la imagen misma, aplanada en un vector de 121 píxeles. coseno aplica la fórmula a dos imágenes aplanadas con flatten. F.cosine_similarity es la versión de PyTorch y calcula la similitud a lo largo de la dimensión dim.

Pregunta. Compara dos vistas de la misma flecha derecha, una centrada y otra desplazada 2 píxeles en filas y columnas, con la flecha izquierda centrada. ¿Qué par tiene mayor similitud coseno en píxeles?

def coseno(a, b):
    a, b = a.flatten(), b.flatten()
    return (a @ b / (a.norm() * b.norm())).item()


vista_1 = lienzo(1)
vista_2 = lienzo(1, 2, 2)
otra = lienzo(0)

print(round(coseno(vista_1, vista_2), 3))
print(round(coseno(vista_1, otra), 3))
print(round(F.cosine_similarity(vista_1.flatten(), vista_2.flatten(), dim=0).item(), 3))

0.2
0.5
0.2

SimCLR aplica esta idea a fotografías. Sus vistas combinan recortes aleatorios, volteo horizontal y cambios de color, y su encoder es una ResNet. Después del entrenamiento sin etiquetas, las representaciones sirven como entrada de un clasificador que se entrena con pocas imágenes etiquetadas. Las transformaciones elegidas deciden qué ignora el encoder. Con fotografías de animales, el volteo horizontal es una vista válida. Con las flechas, el encoder aprendería a ignorar la dirección, que es justo lo que distingue las clases.

Pérdida contrastiva

Una pérdida contrastiva convierte «vistas cercanas, imágenes distintas lejos» en un número que se puede minimizar. Con N imágenes y dos vistas de cada una, hay 2N vectores. Para el vector i:

1. se calcula su similitud coseno con los otros 2N − 1 vectores y se divide entre una temperatura τ
2. el vector de la otra vista de la misma imagen es el positivo; los demás son negativos
3. softmax convierte esas similitudes en probabilidades, y la pérdida del vector es −log de la probabilidad del positivo

La pérdida total es el promedio sobre los 2N vectores. SimCLR llama a esta pérdida NT-Xent.

Con N = 2 imágenes, A y B, y representaciones de dos números ya normalizadas:

Como las cuatro representaciones tienen norma 1, el producto punto de dos de ellas ya es su similitud coseno, y z @ z.T reúne todas las similitudes en una matriz. La celda plegada guarda las cuatro representaciones en z, en el orden A1, A2, B1, B2, y la temperatura en tau.

z = torch.tensor([[1.0, 0.0], [0.8, 0.6], [0.0, 1.0], [-0.6, 0.8]])
tau = 0.5

Pregunta. Con τ = 0.5, calcula las similitudes de A1 con A2, B1 y B2. ¿La probabilidad del positivo de A1 queda por encima o por debajo de 0.5?

Para todos los vectores a la vez, la pérdida es una entropía cruzada. Cada fila de la matriz de similitudes es un vector de logits y la «clase correcta» es el índice del positivo. La diagonal compara cada vector consigo mismo; se reemplaza por -inf para que no participe en el softmax.

F.normalize(..., dim=1) divide cada fila entre su norma, así que z @ z.T contiene cosenos. masked_fill escribe -inf donde la máscara torch.eye vale True, es decir, en la diagonal. Con reduction="none", F.cross_entropy devuelve una pérdida por fila en lugar del promedio.

def perdida_contrastiva(z1, z2, tau):
    n = len(z1)
    z = F.normalize(torch.cat([z1, z2]), dim=1)
    logits = z @ z.T / tau
    logits = logits.masked_fill(torch.eye(2 * n, dtype=torch.bool), float("-inf"))
    # el positivo de la fila i es i + n, y el de la fila i + n es i
    positivos = torch.cat([torch.arange(n, 2 * n), torch.arange(0, n)])
    return F.cross_entropy(logits, positivos, reduction="none")

z1 reúne las primeras vistas, A1 y B1, y z2 las segundas, A2 y B2, en el mismo orden.

Pregunta. ¿Qué vector tiene más pérdida, A1 o A2? mala calcula la pérdida con los mismos cuatro vectores, pero trata A1 y B2 como vistas de la misma imagen, y B1 y A2 como la otra pareja. ¿La pérdida media sube o baja respecto de las parejas originales?

z1 = z[[0, 2]]
z2 = z[[1, 3]]
por_vector = perdida_contrastiva(z1, z2, tau)
mala = perdida_contrastiva(z1, z2[[1, 0]], tau)  # empareja A1 con B2 y B1 con A2

print(por_vector.round(decimals=3))
print(round(por_vector.mean().item(), 3), round(mala.mean().item(), 3))

tensor([0.2330, 0.6270, 0.6270, 0.2330])
0.43 2.03

Práctica

Traza guiada

Antes de procesar, escribe a mano el resultado de torch.flip(t, dims=[1]), torch.flip(t, dims=[0]) y torch.roll(t, shifts=1, dims=1).

t = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])

Comprueba con print. En torch.roll, el 4 de la primera fila debe aparecer en la primera columna. Si no coincide, revisa qué dimensión recorre filas y cuál recorre columnas.

Desplazamiento máximo

Entrena de nuevo con entrenar(lambda x, gen: desplazar(x, gen, maximo=1)) y con maximo=3. Antes de procesar, predice si la exactitud de validación queda por encima o por debajo de la del modelo con maximo=2. Después, cuenta con qué desplazamientos de X_val falla el modelo de maximo=1: guarda dy y dx en conjunto y compáralos con las predicciones.

Tarea con volteo válido

Crea una tarea nueva con las mismas flechas: clase 0 para flechas horizontales (derecha o izquierda) y clase 1 para flechas verticales (arriba o abajo, con torch.rot90). Escribe un cálculo que:

1. genere 40 imágenes centradas de entrenamiento y 200 desplazadas de validación, con las cuatro direcciones en la misma proporción
2. decida, antes de entrenar, qué hacen con la etiqueta el volteo de columnas, el volteo de filas y torch.rot90
3. entrene el MLP de entrenar con las augmentaciones que conservan la etiqueta y compare con una que no la conserva
4. explique, para cada augmentación usada, por qué conserva la etiqueta en esta tarea y no necesariamente en la del problema

Reto opcional

Entrena un encoder sin etiquetas: una capa nn.Conv2d(1, 8, 3), ReLU, un máximo global sobre filas y columnas y nn.Linear(8, 4). En cada paso, toma 16 imágenes de X_ent, genera dos vistas con desplazar y minimiza perdida_contrastiva(...).mean() con τ = 0.5. Compara la similitud coseno media entre vistas de la misma imagen antes y después de 300 pasos. Después, entrena una regresión logística sobre las representaciones con 10 imágenes etiquetadas y mide su exactitud en X_val.

Errores que revisar

• Si la exactitud de validación baja al añadir una augmentación, comprueba si la transformación convierte ejemplos de una clase en ejemplos de otra, como el volteo de una flecha o el giro de 180° de un 6.
• Si el mismo modelo da exactitudes distintas en dos evaluaciones, busca una transformación aleatoria en el código de evaluación.
• Si el volteo de un batch (N, C, H, W) invierte las filas en lugar de las columnas, cambia dims=[2] por dims=[3].
• Si aparecen trazos en el borde opuesto de la imagen después de desplazarla, recuerda que torch.roll es circular; deja un margen vacío de al menos el desplazamiento máximo o rellena con ceros.
• Si la pérdida contrastiva no baja de log 2 ≈ 0.693 aunque las dos vistas sean idénticas, falta enmascarar la diagonal: la similitud de cada vector consigo mismo empata con la de su positivo.

Cierre

Explica por qué el volteo de columnas empeora el clasificador de flechas y por qué tampoco sirve como vista para entrenar un encoder de flechas con pérdida contrastiva.

Lista de comprobación

• En un tensor (N, C, H, W), ¿qué dimensión invierte el volteo horizontal y cuál el vertical?
• ¿Qué tres resultados puede tener una transformación respecto de la etiqueta de una tarea?
• ¿Por qué la augmentación aleatoria se aplica al batch de entrenamiento y no al conjunto de validación?
• ¿Por qué dos vistas de la misma flecha pueden tener menor similitud coseno en píxeles que dos flechas distintas?
• En la pérdida contrastiva, ¿qué vector es el positivo de cada fila y por qué se enmascara la diagonal?

Documentación oficial

• torch.flip: volteo de un tensor en las dimensiones indicadas.
• torch.roll: desplazamiento circular por dimensión.
• torch.nn.functional.cosine_similarity: similitud coseno entre tensores a lo largo de una dimensión.
• Transforming and augmenting images: transformaciones de imágenes de torchvision, como RandomHorizontalFlip y RandomResizedCrop.
• SimCLR: artículo de SimCLR: augmentaciones, encoder ResNet y la pérdida NT-Xent.