Self-supervised обучение в Colab: забираем эмбеддинги из сырых данных

Почему все хотят self-supervised, но никто не делает

Каждый второй пост в ML-сообществе кричит про self-supervised learning. Каждый первый — про то, что разметка данных стоит дорого. Но когда дело доходит до кода, все скатывается к классическому supervised подходу с готовыми датасетами. Почему?

Потому что документация по self-supervised — это академические статьи на 30 страниц с формулами, где код если и есть, то требует кластера из 8 GPU. А в реальности у вас есть Google Colab с его капризным T4 и 12 ГБ памяти, которые исчезают после трех эпох обучения.

Что на самом деле происходит в контрастивном обучении

Представьте, что вы учите ребенка отличать кошек от собак. В supervised подходе вы показываете картинку и говорите: «Это кошка». В self-supervised вы берете фотографию кошки, делаете ее копию, немного меняете (обрезаете, меняют яркость) и говорите: «Эти две картинки должны быть похожи». А потом берете фотографию собаки и говорите: «А эта — не похожа».

Модель учится не классифицировать, а создавать эмбеддинги — компактные представления данных, где похожие объекты находятся близко, а разные — далеко. Эти эмбеддинги потом можно использовать для чего угодно: классификации, кластеризации, поиска похожих объектов.

Косинусная схожесть: почему не евклидово расстояние

Все говорят про cosine similarity, но мало кто понимает, зачем она нужна. Давайте на пальцах.

У вас есть два эмбеддинга: [1, 2, 3] и [2, 4, 6]. Второй в два раза больше первого. По евклидову расстоянию они далеки друг от друга. Но по сути это один и тот же вектор, просто масштабированный. Косинусная схожесть игнорирует длину вектора и смотрит только на направление. Для эмбеддингов это критически важно — нам важно, чтобы похожие объекты лежали в одном направлении, а не обязательно на одинаковом расстоянии от центра.

import torch
import torch.nn.functional as F

# Как НЕ надо делать
def bad_similarity(a, b):
    return torch.sqrt(((a - b) ** 2).sum())  # Евклидово расстояние

# Как надо делать
def good_similarity(a, b):
    return F.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0))  # Косинусная схожесть

# На практике используют нормализацию
embeddings = F.normalize(embeddings, p=2, dim=1)  # L2 нормализация
# После этого косинусная схожесть = скалярное произведение
similarity = torch.mm(embeddings, embeddings.t())

1 Подготовка Colab: что включить, чтобы не плакать потом

Открываете новый ноутбук в Colab. Первое, что делаете — меняете среду выполнения на GPU. Звучит очевидно, но половина людей забывает это сделать и потом удивляется, почему обучение идет 10 часов.

# Проверяем, что GPU доступен
import torch
print(f"GPU доступен: {torch.cuda.is_available()}")
print(f"Название GPU: {torch.cuda.get_device_name(0)}")

# Освобождаем память, если перезапускаем ячейки
torch.cuda.empty_cache()

# Устанавливаем seed для воспроизводимости
import random
import numpy as np

torch.manual_seed(42)
np.random.seed(42)
random.seed(42)

Теперь установите все нужные библиотеки. Не ставьте все подряд — каждая лишняя библиотека съедает память.

!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
!pip install scikit-learn matplotlib tqdm
# Hugging Face datasets понадобится, если работаете с текстом
!pip install datasets transformers

2 Берем данные: где найти и как подготовить

Вам не нужен ImageNet. Возьмите CIFAR-10 или даже CIFAR-100. Маленькие изображения, быстро загружаются, не требуют тонны памяти. Для текста — возьмите какой-нибудь датасет новостей или статей.

Если хотите работать со своими данными, загрузите их на Google Drive и смонтируйте в Colab. Но помните: если у вас 10 000 изображений по 5 МБ каждое, вы быстро упретесь в лимиты.

# Для изображений
from torchvision import datasets, transforms

# Простые аугментации для self-supervised
# Для positive pairs: два случайных кропа от одного изображения
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(32, scale=(0.8, 1.0)),  # Случайный кроп
    transforms.RandomHorizontalFlip(p=0.5),  # Случайное отражение
    transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
    transforms.RandomGrayscale(p=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.247, 0.243, 0.261])
])

# Загружаем CIFAR-10 без меток
train_dataset = datasets.CIFAR10(root='./data', train=True, 
                                 download=True, transform=train_transform)

# Создаем DataLoader
from torch.utils.data import DataLoader
train_loader = DataLoader(train_dataset, batch_size=256, 
                          shuffle=True, num_workers=2, pin_memory=True)

3 Архитектура модели: encoder + projection head

Возьмите готовый ResNet-18, отрежьте последний полносвязный слой — это будет ваш encoder. Добавьте projection head — небольшую MLP, которая превращает эмбеддинги в пространство, где будем считать схожесть.

Зачем projection head? Потому что эмбеддинги, которые хороши для downstream задачи (например, классификации), не обязательно оптимальны для контрастивного обучения. Projection head обучается вместе с моделью, а потом отбрасывается.

import torch.nn as nn
import torchvision.models as models

class SelfSupervisedModel(nn.Module):
    def __init__(self, feature_dim=512, projection_dim=128):
        super().__init__()
        
        # Encoder - берем ResNet без головы
        backbone = models.resnet18(pretrained=False)  # Можно True, если хотите transfer learning
        self.encoder = nn.Sequential(*list(backbone.children())[:-1])  # Убираем avgpool и fc
        
        # Projection head
        self.projector = nn.Sequential(
            nn.Linear(512, feature_dim),  # ResNet-18 выдает 512 фичей
            nn.BatchNorm1d(feature_dim),
            nn.ReLU(),
            nn.Linear(feature_dim, projection_dim)
        )
        
    def forward(self, x):
        # Получаем фичи от encoder
        features = self.encoder(x)
        features = features.view(features.size(0), -1)  # Flatten
        
        # Проецируем в пространство для контрастивного обучения
        projection = self.projector(features)
        
        # Возвращаем и фичи (для downstream задач), и проекцию
        return features, projection

4 Контрастивный лосс: NT-Xent на практике

Normalized Temperature-scaled Cross Entropy Loss — страшное название для простой идеи. Для каждого изображения в батче у вас есть positive pair (два аугментированных варианта) и 2*(batch_size-1) negative примеров (все остальные изображения в батче).

Лосс штрафует модель, если positive pair недостаточно похожи, и если negative pair слишком похожи.

class NTXentLoss(nn.Module):
    def __init__(self, temperature=0.5):
        super().__init__()
        self.temperature = temperature
        self.cosine_similarity = nn.CosineSimilarity(dim=2)
        
    def forward(self, projections):
        # projections: [2*batch_size, projection_dim]
        # Первые batch_size - оригинальные изображения
        # Вторые batch_size - аугментированные версии
        
        batch_size = projections.size(0) // 2
        
        # Нормализуем проекции
        projections = F.normalize(projections, p=2, dim=1)
        
        # Вычисляем матрицу схожести
        similarity_matrix = torch.mm(projections, projections.t()) / self.temperature
        
        # Маска для positive pairs
        pos_mask = torch.zeros_like(similarity_matrix)
        pos_mask[:batch_size, batch_size:2*batch_size] = 1
        pos_mask[batch_size:2*batch_size, :batch_size] = 1
        
        # Маска для negative pairs (исключаем диагональ)
        neg_mask = 1 - pos_mask
        neg_mask.fill_diagonal_(0)
        
        # Лосс для positive pairs
        pos_similarities = similarity_matrix * pos_mask
        pos_loss = -pos_similarities.sum() / (2 * batch_size)
        
        # Лосс для negative pairs (logsumexp trick для численной стабильности)
        neg_similarities = similarity_matrix * neg_mask
        neg_loss = torch.logsumexp(neg_similarities, dim=1).sum() / (2 * batch_size)
        
        return pos_loss + neg_loss

5 Обучение: как не перегреть Colab

Самая частая ошибка — пытаться обучить ResNet-50 на батче 512. Colab с T4 просто не потянет. Используйте ResNet-18 или даже ResNet-9. Батч 128-256 — оптимальный вариант.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SelfSupervisedModel().to(device)
criterion = NTXentLoss(temperature=0.5)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Learning rate scheduler
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

# Тренировочный цикл
model.train()
for epoch in range(50):
    total_loss = 0
    
    for batch_idx, (images, _) in enumerate(train_loader):  # Метки не используем
        images = images.to(device)
        
        # Создаем два аугментированных представления
        # В реальности нужно применять разные аугментации
        images_aug1 = images  # Уже аугментированы в transform
        images_aug2 = images  # Для простоты используем те же
        
        # Конкатенируем для одного прохода
        batch = torch.cat([images_aug1, images_aug2], dim=0)
        
        # Forward pass
        _, projections = model(batch)
        loss = criterion(projections)
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        
        # Мониторинг памяти
        if batch_idx % 50 == 0:
            print(f"Batch {batch_idx}, Loss: {loss.item():.4f}")
            print(f"GPU memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
    
    scheduler.step()
    print(f"Epoch {epoch}, Average Loss: {total_loss / len(train_loader):.4f}")

Что делать с эмбеддингами после обучения

Вы обучили модель. У вас есть encoder, который превращает изображения в 512-мерные векторы. Что дальше?

1. Кластеризация — используйте K-means или DBSCAN на эмбеддингах, чтобы найти группы похожих изображений без меток.
2. Поиск похожих — по косинусной схожести находите самые близкие эмбеддинги в базе.
3. Transfer learning — заморозьте encoder, добавьте линейный слой и дообучите на маленьком размеченном датасете. Это работает лучше, чем обучение с нуля.
4. Визуализация — с помощью t-SNE или UMAP превратите 512-мерные векторы в 2D и посмотрите, как данные распределяются.

# Пример: поиск похожих изображений
model.eval()
with torch.no_grad():
    # Получаем эмбеддинги для всех изображений
    all_embeddings = []
    for images, _ in train_loader:
        images = images.to(device)
        embeddings, _ = model(images)
        all_embeddings.append(embeddings.cpu())
    
    all_embeddings = torch.cat(all_embeddings, dim=0)
    all_embeddings = F.normalize(all_embeddings, p=2, dim=1)
    
    # Берем тестовое изображение
    test_idx = 0
    test_embedding = all_embeddings[test_idx]
    
    # Ищем похожие
    similarities = torch.mm(all_embeddings, test_embedding.unsqueeze(1)).squeeze()
    top_indices = similarities.argsort(descending=True)[:5]
    
    print(f"Самые похожие изображения: {top_indices.tolist()}")

Где все ломается: частые ошибки и как их исправить

Self-supervised vs другие подходы

Если вам нужно быстро получить эмбеддинги для текста, посмотрите RAG за 15 минут. Там есть готовые примеры работы с эмбеддингами.

Для тонкой настройки LLM есть полное руководство.

Вместо эпилога: готовый Colab-ноутбук

Весь код из статьи, включая обучение на CIFAR-10, визуализацию эмбеддингов и поиск похожих изображений, доступен в виде ноутбука. Вы можете запустить его прямо в Colab и посмотреть, как работает self-supervised learning на практике.

Не забудьте про GPU и свободную память. Удачи!