solution

README.md

@@ -10,3 +10,226 @@ pinned: false
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+
+# Cross Attention Classifier
+
+Ниже технические детали того, как устроен репозиторий и как обучить модель.
+В самой последней секции "Описание подходов" подробно описано, как я пришла к этому подходу, с какими проблемами встретилась, а также описаны два других подхода, которые я решила не реализовывать (будет мини-эссе, готовьтесь).
+
+## Описание проекта
+
+В проекте используется self-supervised обучение (BYOL) и последующая классификацию изображений с помощью Cross Attention.  
+- **train_byol.py** — скрипт для обучения модели-энкодера по методу BYOL.  
+- **train_cross_classifier.py** — скрипт для обучения классификатора, который использует предварительно обученный энкодер и Cross Attention.  
+- **app.py** — Streamlit-приложение для инференса и визуализации предсказаний (генерация случайных изображений и получение метки от модели).
+
+## Структура репозитория
+
+```
+.
+├── src
+│   ├── dataset.py            # Реализация датасетов (RandomAugmentedDataset и RandomPairDataset)
+│   ├── inference.py          # Класс для инференса (CrossAttentionInference) и вспомогательные методы
+│   └── models.py             # Определения моделей (BYOL, VGGLikeEncode, CrossAttentionClassifier)
+├── train_byol.py             # Скрипт обучения модели BYOL
+├── train_cross_classifier.py # Скрипт обучения Cross Attention Classifier (использует готовый энкодер)
+├── app.py                    # Streamlit-приложение для инференса
+├── requirements.txt          # Список Python-зависимостей (pip install -r requirements.txt)
+└── pyproject.toml / poetry.lock # Файл для установки зависимостей через Poetry
+```
+
+## Установка зависимостей
+
+Можно установить зависимости двумя способами:  
+1. **Через `pip` и `requirements.txt`:**  
+   ```bash
+   pip install -r requirements.txt
+   ```
+2. **Через Poetry:**  
+   ```bash
+   poetry install
+   ```
+
+## Как обучить модель
+
+### 1. Обучение энкодера с помощью BYOL
+Нужно запустить:
+```bash
+python train_byol.py
+```
+- Этот скрипт обучает модель энкодера (`VGGLikeEncode`) методом BYOL на данных, сгенерированных `RandomAugmentedDataset`.
+- После обучения лучшая модель (с минимальным `val_loss`) сохраняется в `best_byol.pth`.
+
+### 2. Обучение Cross Attention Classifier
+
+`best_byol.pth` (веса энкодера) должны лежать в корневой папке (можно указать другой путь). Затем нужно запустить:
+```bash
+python train_cross_classifier.py
+```
+- Этот скрипт использует предобученный энкодер и обучает классификатор для определения, содержат ли картинки одинаковую геометрическую фигуру.
+- По итогам сохранит веса модели-классификатора в `best_attention_classifier.pth`.
+
+## Как запустить инференс
+
+### Запуск через Streamlit-приложение
+
+1. Файл весов `best_attention_classifier.pth` должен лежать в корневой папке
+2. Нужно запустить Streamlit-приложение:
+   ```bash
+   streamlit run app.py
+   ```
+3. Дальше, нужно перейти по адресу, который выдаст Streamlit (по умолчанию [http://localhost:8501](http://localhost:8501)).  
+4. Нажмите кнопку **«Сгенерировать изображения»**. Приложение сгенерирует пару случайных изображений и покажет предсказанную моделью метку.
+
+### Использование класса инференса в коде напрямую
+
+Можно использовать модель напрямую (без интерфейса Streamlit), импортируйте класс из `src/inference.py`, передайте путь к весам модели и вызовите метод предсказания. Пример:
+
+```python
+import torch
+from src.inference import CrossAttentionInference
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+inference = CrossAttentionInference(
+    model_path="best_attention_classifier.pth",
+    device=device
+)
+
+pred_label, (img1, img2) = inference.predict_random_pair()
+print(f"Предсказанная метка: {pred_label}")
+```
+
+---
+
+## Описание подходов
+
+
+### BYOL + Cross-attention (выбранный подход)
+
+Когда я обдумывала финальную архитектуру, я поняла, что у креативности нет предела, поэтому каждое решение должно быть обосновано не только фразой "прикольно, можно попробовать", но и существующими проблемами, которые хочется решить.
+И я решила отталкиваться от реальных задач, а именно от проблемы с отсутствием данных в медицинской сфере.
+В текущей задаче такой проблемы, очевидно, нет, мы можем сгенерировать хоть миллион изображений и для всех будет лейбл.
+Но вот что, если у нас нет возможности сгенерировать миллион изображений? Или если у нас есть 100к изображений, но только 10000 из них размечены? 
+Например, у нас есть неплохой банк изображений с разными опухолями, но размечены только 10% из них. Как можно использовать эти данные для обучения модели, чтобы она могла классифицировать новые изображения?
+Саму же задачу можно перенести на задачу вида "мониторинг прогрессирования заболевания" или "сравнение патологий".
+
+Поэтому я решила использовать self-supervised обучение для того, чтобы обучить модель на неразмеченных данных и затем дообучить ее на небольшом датасете с разметкой.
+
+#### BYOL
+BYOL [Bootstrap Your Own Latent](https://arxiv.org/pdf/2006.07733) — это метод self-supervised обучения, который позволяет обучить модель на неразмеченных данных. 
+Важной особенностью конкретно этого подхода заключается в том, что этот метод не требует негативных пар для обучения, как некоторые другие contrastive методы. 
+В BYOL используется две копии одной и той же модели, которые обучаются предсказывать друг друга на основе двух views (аугментаций) одного изображения. 
+Не схлапываться в один вектор помогает то, что архитектура не симметрична, так как в одной из веток добавляется MLP предиктор, а также stop gradient операция.
+В итоге модель учится извлекать признаки из изображения, которые можно использовать для дообучения на меньшем датасете.  
+В данном случае я использовала энкодер с похожей на `VGG` архитектурой.
+Выбрала VGG я потому, что использовать, например, ResNet со skip-connection нет смысла, так как изображение всего ли 32x32, и через несколько слоев feature мапа была уже 8x8.
+
+В целом сама задача заставляет балансировать между сложными подходами и реальной возможностью обучить модель на подобном датасете, так как реализовать можно (почти) что угодно, но переобучиться на таком датасете достаточно легко.
+
+Изначально я планировала добавить в датасет для предобучения другие фигуры (треугольники, звездочки и тд), но сами эти фигуры занимают несколько пикселей, и аугментации их сильно искаж��ют. В целом на таких маленьких изображениях почти все аугментации становятся агрессивными.
+Поэтому я остановилась на двух фигурах, но добавила в аугментации реверс цвета, повороты, сдвигы, гауссовский шум и тд.
+
+#### Cross Attention
+Честно скажу, именно на этот подход меня вдохновила задача с собеседования, где нужно было сопоставить два снимка одной и той же области.
+Я нашла статью - [An Adaptive Remote Sensing Image-Matching Network Based on Cross Attention and Deformable Convolution](https://www.researchgate.net/publication/388063503_An_Adaptive_Remote_Sensing_Image-Matching_Network_Based_on_Cross_Attention_and_Deformable_Convolution)
+где авторы решают похожую задачу (они тоже кстати используют VGG), но для более сложных изображений (сопоставление фотографий со спутника).
+Я помню в чем заключается проблема cross-attention - у него квадратичная сложность, и если изображение имеет размер 512x512, то это уже становится проблемой. 
+Но так как в задаче изображения 32x32, я решила, что будет уместно применить данный подход (в предложенных подходах дальше я опишу, как бы решала задачу, если бы изображения были больше).
+Также, я добавила position эмбеддинги, так как при переходе к cross-attention информация о позиции теряется.
+
+Почему cross-attention? Он позволяет каждому пикселю (точнее, каждому патчу) в одном изображении "смотреть" на все патчи в другом изображении.
+Таким образом, если фигуры находятся в противоположных углах, модель это учтет.
+Ну и плюс тенденции последних лет - внимание, внимание, внимание.
+
+#### Итоговая архитектура
+Сама архитектура представляет собой два VGGLike энкодера с shared весами, предобученных с помощью BYOL, после которых идет слой MultiheadAttention, а затем классификационная голова.
+Во время предобучения VGGLike энкодера последним слоем был AdaptiveAvgPool2d. Этот слой не использовался во время обучения классфикатора, так как на вход MultiheadAttention требовалась информативная карта признаков (я использовала 8x8).
+
+Таким образом, когда на вход поступает два изображения, каждое из них проходит через энкодер, после чего происходит cross-attention между ними, и на выходе получается вероятность того, что изображения содержат одинаковую фигуру.
+Это не самый сложный подход, который можно было придумать, но он позволяет взглянуть на задачу под другим углом - в реальности у нас нет датасета с неограниченным количеством размеченных данных, и нужно уметь работать с тем, что есть.
+
+### Метрики
+
+Вот здесь можно посмотреть метрики в wandb:
+- [Обучение BYOL](https://wandb.ai/alexandraroze/contrastive_learning_byol/reports/-BYOL--VmlldzoxMTQzMjA1Mw?accessToken=nh0kzpepsr0faflptx63n91kljc5wl6mt3wi3ay4wxpjmua55bf32nm36qjby0ai)
+- [Обучение cross-attention классификатора](https://api.wandb.ai/links/alexandraroze/hmtnzhv9)
+
+
+## Другие подходы
+
+
+### Swin transformer
+
+
+1. Каждое 32×32 изображение делится на патчи размером 4×4, это даёт 64 патча на изображение.
+Каждый патч выпрямляется и проходит через линейный слой для получения векторного представления.
+
+2. Далее мы применим early fusion (так как если применить late fusion, нам придется применять cross-attention, чтобы действительно учест�� взаимодействие между патчами из разных изображений).
+После извлечения патчей из двух изображений мы просто конкатенируем их по оси последовательности, получая 128 токенов.
+
+3. В window multi-head attention мы делим эту последовательность на окна фиксированного размера. Допустим, каждое окно включает 16 токенов подряд. Это значит, что фигура, находящаяся в определённом блоке патчей, будет анализироваться локально вместе со смежными патчами.
+Применяем self-attention и затем сдвигаем окна (в целом, как и должно быть в swin blockе).
+Дальше идет patch merging, и мы получаем 16 патчей на одно изображение (то есть 32 патча на два изображения).
+Достаточно еще двух таких слоев (16 -> 4 -> 1), чтобы у нас остался один патч на изображение.
+
+4. Далее мы используем global average pooling, и передаем выход в классификационную голову.
+
+#### Почему я не стала реализовывать этот подход
+Swin transformer хорошо сработает на крупных изображениях с мелкими деталями, но в данной задаче спустя всего 3 слоя мы уже получаем один токен на изображение.
+В первом слое локальное внимание ограничено работает сразу для двух изображений, а в следующем слое остается уже не так много токенов, чтобы извлекать информацию о фигурах.
+
+
+
+### Siamese Network с Triplet Loss
+
+Вместо простой классификации мы обучаем энкодер, который преобразует изображения в эмбеддинги так, чтобы похожие изображения (круг-круг, квадрат-квадрат) были ближе друг к другу, а разные изображения (круг-квадрат) были дальше.
+
+Используем Triplet Loss, где берём три изображения:
+- Anchor – произвольное изображение (например, квадрат).
+- Positive – ещё один квадрат.
+- Negative – круг.
+
+Модель минимизирует расстояние между anchor и positive и максимизирует его для negative. Чтобы модели было сложнее, используем hard negatives. Например, генерировать изображение с одинаковыми характеристиками, такими как положение фигуры, цвет, блюр, но с другой фигурой.
+
+Используем легкий shared CNN энкодер. Энкодер обрабатывает изображения независимо, но выходные эмбеддинги сравниваются через triplet loss.
+Важно, чтобы размерность эмбеддинга была достаточно низкой, чтобы не переобучиться на простой структуре.
+
+Получаем эмбеддинги двух изображений, считаем евклидово расстояние. 
+Если меньше порога - фигуры одинаковые, иначе разные.
+
+
+#### Почему я не стала реализовывать этот подход
+Я уже делала это на своей работе, поэтому хотелось попробовать что-нибудь новое :)
+
+
+### Проблемы с которыми я столкнулась
+1. Вначале я решила написать и обучить всю архитектуру целиком (энкодер + cross-attention классификатор), но сразу же столкнулась с тем, что модель просто не обучалась.
+   Чтобы это отдебажить, я решила начать с малого - создала простой датасет и научила простую CNN предсказывать метку для двух изображений сразу.
+   Дальше, я добавляла углубление в энкодер, параллельно мониторя количество параметров, чтобы понимать, какое количество сэмплов мне нужно для обучения.
+   Та��им образом, я дошла до финальной архитектуры. 
+2. У меня все еще сохранялись проблемы во время обучения (сеть не обучалась). Мониторинг нормы градиентов и весов помог мне понять, что веса из-за attention просто зануляются. Это я решила изменением оптимизатора на AdamW и уменьшением learning rate. 
+3. Изначально планировалось показать улучшения в обучении с помощью self-supervised обучения (в сравнении с обучением с нуля), но по факту при тех же самых условиях обучение проходило одинаково в обоих случаях.
+Это можно объяснить тем, что изображения были слишком маленькими и простыми, и подходу без предобучения также не требовалось много времени и большого количества данных. 
+Чтобы self-supervised метод действительно хорошо сработал (особенно без негативных примеров), нужны сложные аугментации, в этом же случае сложные аугментации сильно искажали изображения.
+4. Судя по кривой обучения, итоговый классификатор очень долгое время находился на плато, так как первые 7-8 эпох из 10 лосс не падал, а точность оставалась на уровне 50%. Это можно объяснить тем, что градиенты очень маленькие или очень шумные. Также, все зависит от исходной инициализации, и при маленьком датасете это может стать проблемой, так как по началу накапливается недостаточно сигналов, чтобы сойти с плато.
+
+
+## Что бы я точно не стала делать
+Здесь я опишу подходы, которые сразу пришли мне в голову, но которые я бы точно не стала делать по итогу.
+Опять же, я отталкивалась от переноса задачи на реальные данные.
+
+1. Сверточная сеть, которая принимает на вход одно изображение, и выдает для него класс (круг или квадрат). 
+Соответственно, получив предсказания для двух изображений, мы можем сделать вывод о том, содержат ли они одинаковую фигуру.
+Этот подход очень простой, решает задачу в лоб, но он не масштабируется, так как при переносе на реальные кейсы терпит крах, потому что далеко не всегда у нас есть два четко разделенных класса (да и в целом у нас может и не быть классов, а только изображения, которые нужно сопоставить между собой).
+
+2. Детекция + классификация.
+Можно было достаточно просто обучить детектор, который находил бы как определенный класс (круг или квадрат), так и просто "фигуру" без класса (казалось бы, решение предыдущей проблемы).
+В реальности же этот подход тоже не масштабируется, так как 1) это дорогостоящая разметка, 2) детекторы могут ошибаться (и для таких кейсов мы бы тогда вообще могли ничего не предсказать), 3) задача может состоять в сравнении нескольких разнородных объектов на изображении, а не одного (например, образование новых опухолей).
+То же само касается и сегментации.
+
+
+Здесь стоит сделать важную поправку, что есть реальные задачи, где эти подходы могут сработать (например, мы точно знаем, что на изображении нас интересует только один объект, а все остальное - неинформативный фон).
+
+

app.py

@@ -0,0 +1,30 @@
+import streamlit as st
+import torch
+
+from src.inference import CrossAttentionInference
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+inference = CrossAttentionInference(
+    model_path="best_attention_classifier.pth",
+    device=device
+)
+
+st.title("Random Image Inference")
+
+st.write(
+    "Нажмите кнопку ниже, чтобы сгенерировать пару случайных изображений и получить предсказание модели."
+)
+
+if st.button("Сгенерировать изображения"):
+    pred_label, (img1, img2) = inference.predict_random_pair()
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        st.image(img1, caption="Image 1", use_container_width=True)
+    with col2:
+        st.image(img2, caption="Image 2", use_container_width=True)
+
+    st.write(f"**Предсказанная метка**: {pred_label}")

best_attention_classifier.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6496c9cd313964eef9b899924b8429b11b257cf6ea78a18af3a06df2ed16afb8
+size 1527896

best_byol.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a09baac108efe3034c6f47aa174ba7afb3e8b2732aad4598f9e24ca642337d0
+size 1158467

pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "cels-test"
+version = "0.1.0"
+description = ""
+authors = [
+    {name = "AleksandraSorokovikova",email = "alexandraroze2000@gmail.com"}
+]
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "torch (>=2.6.0,<3.0.0)",
+    "torchvision (>=0.21.0,<0.22.0)",
+    "matplotlib (>=3.10.0,<4.0.0)",
+    "wandb (>=0.19.6,<0.20.0)",
+    "tqdm (>=4.67.1,<5.0.0)",
+    "streamlit (>=1.42.0,<2.0.0)"
+]
+
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"

src/dataset.py

@@ -0,0 +1,236 @@
+import random
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torchvision.transforms as T
+from PIL import Image, ImageDraw, ImageFilter
+from torch.utils.data import Dataset
+
+
+def generate_image(
+        size: int = 32,
+        contrast: Tuple[int, int] = (90, 110),
+        blur_radius: Tuple[float, float] = (0.5, 1.5),
+        shape: Optional[str] = None,
+        max_background_intensity: int = 128,
+        min_shape_intensity: Optional[int] = None,
+        shape_size: Optional[int] = None,
+        location: str = 'random',
+        random_intensity: bool = False
+) -> Tuple[Image.Image, str]:
+    """
+    Generate an image with a shape (circle or square) on a background.
+    :param size: size of the image
+    :param contrast: contrast of the shape
+    :param blur_radius: radius of the Gaussian blur
+    :param shape: shape type (circle or square)
+    :param max_background_intensity: maximum intensity of the background
+    :param min_shape_intensity: minimum intensity of the shape
+    :param shape_size: size of the shape
+    :param location: location of the shape ('random' or 'center')
+    :param random_intensity: whether to randomly invert the shape intensity
+    """
+    background_intensity = random.randint(0, max_background_intensity)
+    background = Image.new('L', (size, size), background_intensity)
+
+    if shape:
+        assert shape in ['circle', 'square'], "Wrong shape type"
+    else:
+        shape = random.choice(['circle', 'square'])
+
+    if not min_shape_intensity:
+        random_contrast = random.randint(*contrast)
+        min_shape_intensity = min(background_intensity + random_contrast, 255)
+    shape_intensity = random.randint(min_shape_intensity, 255)
+
+    mask = Image.new('L', (size, size), 0)
+    draw = ImageDraw.Draw(mask)
+
+    if not shape_size:
+        min_size = 8
+        max_size = size // 2
+        shape_size = random.randint(min_size, max_size)
+
+    if location == 'random':
+        max_pos = size - shape_size - 1
+        top_left_x = random.randint(0, max_pos)
+        top_left_y = random.randint(0, max_pos)
+    else:
+        top_left_x = (size - shape_size) // 2
+        top_left_y = (size - shape_size) // 2
+
+    if shape == 'square':
+        draw.rectangle([top_left_x, top_left_y, top_left_x + shape_size, top_left_y + shape_size], fill=255)
+    else:
+        draw.ellipse([top_left_x, top_left_y, top_left_x + shape_size, top_left_y + shape_size], fill=255)
+
+    if blur_radius:
+        random_blur_radius = random.uniform(*blur_radius)
+        mask = mask.filter(ImageFilter.GaussianBlur(radius=random_blur_radius))
+    else:
+        mask = mask.filter(ImageFilter.SMOOTH)
+
+    shape_img = Image.new('L', (size, size), shape_intensity)
+    img = Image.composite(shape_img, background, mask)
+
+    if random_intensity and random.random() < 0.5:
+        img = Image.eval(img, lambda x: 255 - x)
+
+    return img, shape
+
+
+class RandomPairDataset(Dataset):
+    def __init__(
+            self,
+            shape_params: Optional[dict] = None,
+            num_samples: int = 1000,
+            train: bool = True,
+            fixed_test_data: Optional[list] = None
+    ):
+        """
+        Dataset for training a model to compare two images.
+        :param shape_params: parameters for generate_image function
+        :param num_samples: number of samples in the dataset
+        :param train: whether to generate training or test data
+        :param fixed_test_data: fixed test data (optional)
+        """
+        self.train = train
+        self.num_samples = num_samples
+        self.transform = T.Compose([
+            T.ToTensor(),
+            T.Normalize(mean=(0.5,), std=(0.5,))
+        ])
+        if not shape_params:
+            self.shape_params = {}
+        else:
+            self.shape_params = shape_params
+
+        if not self.train:
+            if fixed_test_data is None:
+                self.data = [self._generate_pair() for _ in range(num_samples)]
+            else:
+                self.data = fixed_test_data
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self.train:
+            img1, shape1, img2, shape2, label = self._generate_pair()
+        else:
+            img1, shape1, img2, shape2, label = self.data[idx]
+
+        img1 = self.transform(img1)
+        img2 = self.transform(img2)
+
+        return img1, img2, torch.tensor(label, dtype=torch.float32)
+
+    def _generate_pair(self) -> Tuple[Image.Image, str, Image.Image, str, int]:
+        img1, shape1 = generate_image(**self.shape_params)
+        img2, shape2 = generate_image(**self.shape_params)
+        label = 1 if shape1 == shape2 else 0
+
+        return img1, shape1, img2, shape2, label
+
+
+class RandomAugmentedDataset(Dataset):
+    def __init__(
+            self,
+            augmentations: T.Compose,
+            shape_params: Optional[dict] = None,
+            num_samples: int = 1000,
+            train: bool = True,
+            fixed_test_data: Optional[list] = None
+    ):
+        """
+        Dataset for training a model with contrastive learning.
+        :param augmentations: augmentations to apply to the images
+        :param shape_params: parameters for generate_image function
+        :param num_samples: number of samples in the dataset
+        :param train: whether to generate training or test data
+        :param fixed_test_data: fixed test data (optional
+        """
+        self.train = train
+        self.num_samples = num_samples
+        self.augmentations = augmentations
+        if not shape_params:
+            self.shape_params = {}
+        else:
+            self.shape_params = shape_params
+
+        if not self.train:
+            if fixed_test_data is None:
+                self.data = [self._generate_single() for _ in range(num_samples)]
+            else:
+                self.data = fixed_test_data
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.train:
+            img, _ = self._generate_single()
+        else:
+            img, _ = self.data[idx]
+        view_1, view_2 = self.augmentations(img), self.augmentations(img)
+
+        return view_1, view_2
+
+    def _generate_single(self) -> Tuple[Image.Image, int]:
+        img, shape = generate_image(**self.shape_params)
+        label = 1 if shape == "circle" else 0
+
+        return img, label
+
+
+class AddGaussianNoise(object):
+    def __init__(self, mean: float = 0.0, std: float = 0.05):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, tensor: torch.Tensor) -> torch.Tensor:
+        noise = torch.randn(tensor.size()) * self.std + self.mean
+        tensor = tensor + noise
+        return torch.clamp(tensor, 0., 1.)
+
+    def __repr__(self):
+        return f'{self.__class__.__name__}(mean={self.mean}, std={self.std})'
+
+
+class ColorInversion(object):
+    def __call__(self, image: Image.Image) -> Image.Image:
+        return Image.eval(image, lambda x: 255 - x)
+
+    def __repr__(self):
+        return f'{self.__class__.__name__}()'
+
+
+def get_byol_transforms() -> T.Compose:
+    """
+    Get augmentations for training with BYOL.
+    """
+    augmentations = T.Compose([
+        T.RandomResizedCrop(size=32, scale=(0.8, 1.0), ratio=(0.9, 1.1)),
+        T.RandomHorizontalFlip(p=0.5),
+        T.RandomVerticalFlip(p=0.5),
+        T.RandomRotation(degrees=15),
+        T.ColorJitter(brightness=0.2, contrast=0.2),
+        T.RandomApply([T.GaussianBlur(kernel_size=3, sigma=(0.1, 1.0))], p=0.5),
+        T.RandomApply([ColorInversion()]),
+        T.ToTensor(),
+        T.RandomApply([AddGaussianNoise(mean=0.0, std=0.05)], p=0.5),
+        T.Normalize(mean=(0.5,), std=(0.5,))
+    ])
+    return augmentations
+
+
+def tensor_to_image(tensor: torch.Tensor) -> Image.Image:
+    """
+    Convert a tensor to a PIL image.
+    """
+    img_norm = tensor.cpu()[0]
+    img_denorm = img_norm * 0.5 + 0.5
+    arr = (img_denorm.numpy() * 255).astype(np.uint8)
+    pil_img = Image.fromarray(arr, mode='L')
+    return pil_img

src/inference.py

@@ -0,0 +1,59 @@
+import torch
+import torchvision.transforms as T
+
+from typing import Optional
+
+from src.dataset import generate_image
+from src.models import CrossAttentionClassifier, VGGLikeEncode
+
+
+class CrossAttentionInference:
+    def __init__(
+        self,
+        model_path: str,
+        shape_params: Optional[dict] = None,
+        device: torch.device = torch.device("cpu"),
+    ):
+        if not shape_params:
+            self.shape_params = {}
+        else:
+            self.shape_params = shape_params
+        self.device = device
+
+        self.encoder = VGGLikeEncode(
+            in_channels=1,
+            out_channels=128,
+            feature_dim=32,
+            apply_pooling=False
+        )
+        self.model = CrossAttentionClassifier(encoder=self.encoder)
+
+        state_dict = torch.load(model_path, map_location=device)
+        self.model.load_state_dict(state_dict)
+
+        self.model.eval()
+        self.model.to(device)
+
+        self.transform = T.Compose([
+            T.ToTensor(),
+            T.Normalize(mean=(0.5,), std=(0.5,))
+        ])
+
+    def pil_to_tensor(self, img):
+        return self.transform(img).unsqueeze(0).to(self.device)
+
+
+    def predict_random_pair(self):
+        img1, _ = generate_image(**self.shape_params)
+        img2, _ = generate_image(**self.shape_params)
+
+        img1_tensor = self.pil_to_tensor(img1)
+        img2_tensor = self.pil_to_tensor(img2)
+
+        with torch.no_grad():
+            logits, _ = self.model(img1_tensor, img2_tensor)
+
+        preds = (torch.sigmoid(logits) > 0.5).float()
+        predicted_label = int(preds.item())
+
+        return predicted_label, (img1, img2)

src/models.py

@@ -0,0 +1,241 @@
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class VGGLikeEncode(nn.Module):
+    def __init__(
+            self,
+            in_channels: int = 1,
+            out_channels: int = 128,
+            feature_dim: int = 32,
+            apply_pooling: bool = False
+    ):
+        """
+        VGG-like encoder for grayscale images.
+        :param in_channels: number of input channels
+        :param out_channels: number of output channels
+        :param feature_dim: number of channels in the intermediate layers
+        :param apply_pooling: whether to apply global average pooling at the end
+        """
+        super().__init__()
+        self.apply_pooling = apply_pooling
+
+        self.block1 = nn.Sequential(
+            nn.Conv2d(in_channels, feature_dim, kernel_size=3, padding=1),
+            nn.BatchNorm2d(feature_dim),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(feature_dim, feature_dim, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2)
+        )
+
+        self.block2 = nn.Sequential(
+            nn.Conv2d(feature_dim, feature_dim * 2, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm2d(feature_dim * 2),
+            nn.Conv2d(feature_dim * 2, feature_dim * 2, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2)
+        )
+
+        self.block3 = nn.Sequential(
+            nn.Conv2d(feature_dim * 2, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm2d(out_channels),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=1)
+        )
+
+        self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.blocks = [self.block1, self.block2, self.block3]
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.block1(x)
+        x = self.block2(x)
+        x = self.block3(x)
+        if self.apply_pooling:
+            x = self.global_avg_pool(x).view(x.shape[0], -1)
+        return x
+
+    def get_conv_layer(self, block_num: int):
+        if block_num >= len(self.blocks):
+            return None
+        return self.blocks[block_num][0]
+
+
+class CrossAttentionClassifier(nn.Module):
+    def __init__(
+            self,
+            feature_dim: int = 32,
+            num_heads: int = 4,
+            linear_dim: int = 128,
+            out_channels: int = 128,
+            encoder: Optional[VGGLikeEncode] = None
+    ):
+        """
+        Cross-attention classifier for comparing two grayscale images.
+        :param feature_dim: number of channels in the intermediate layers
+        :param num_heads: number of attention heads
+        :param linear_dim: number of units in the linear layer
+        :param out_channels: number of output channels
+        :param encoder: encoder to use
+        """
+        super(CrossAttentionClassifier, self).__init__()
+        if encoder:
+            self.encoder = encoder
+        else:
+            self.encoder = VGGLikeEncode(in_channels=1, feature_dim=feature_dim, out_channels=out_channels)
+
+        self.out_channels = out_channels
+        self.seq_len = 8 * 8
+        self.pos_embedding = nn.Parameter(torch.randn(self.seq_len, 1, out_channels) * 0.01)
+
+        self.cross_attention = nn.MultiheadAttention(
+            embed_dim=out_channels,
+            num_heads=num_heads,
+            batch_first=False
+        )
+
+        self.norm = nn.LayerNorm(out_channels)
+
+        self.classifier = nn.Sequential(
+            nn.Linear(out_channels, linear_dim),
+            nn.ReLU(),
+            nn.Linear(linear_dim, 1)
+        )
+
+    def forward(self, img1: Tensor, img2: Tensor) -> Tuple[Tensor, Tensor]:
+        feat1 = self.encoder(img1)
+        feat2 = self.encoder(img2)
+
+        B, C, H, W = feat1.shape
+        seq_len = H * W
+
+        feat1_flat = feat1.view(B, C, seq_len).permute(2, 0, 1)
+        feat2_flat = feat2.view(B, C, seq_len).permute(2, 0, 1)
+
+        feat1_flat = feat1_flat + self.pos_embedding
+        feat2_flat = feat2_flat + self.pos_embedding
+
+        feat1_flat = self.norm(feat1_flat)
+        feat2_flat = self.norm(feat2_flat)
+
+        attn_output, attn_weights = self.cross_attention(
+            query=feat1_flat,
+            key=feat2_flat,
+            value=feat2_flat,
+            need_weights=True,
+            average_attn_weights=True
+        )
+        pooled_features = attn_output.mean(dim=0)
+        logits = self.classifier(pooled_features).squeeze(-1)
+
+        return logits, attn_weights
+
+
+class NormalizedMSELoss(nn.Module):
+    def __init__(self):
+        """
+        Normalized MSE loss for BYOL training.
+        """
+        super(NormalizedMSELoss, self).__init__()
+
+    def forward(self, view1: Tensor, view2: Tensor) -> Tensor:
+        v1 = F.normalize(view1, dim=-1)
+        v2 = F.normalize(view2, dim=-1)
+        return 2 - 2 * (v1 * v2).sum(dim=-1)
+
+
+class MLP(nn.Module):
+    def __init__(self, input_dim: int, projection_dim: int = 128, hidden_dim: int = 512):
+        """
+        MLP for BYOL training.
+        :param input_dim: input dimension
+        :param projection_dim: projection dimension
+        :param hidden_dim: hidden dimension
+        """
+        super(MLP, self).__init__()
+
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.BatchNorm1d(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim, projection_dim)
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.net(x)
+
+
+class EncoderProjecter(nn.Module):
+    def __init__(self, encoder: nn.Module, hidden_dim: int = 512, projection_out_dim: int = 128):
+        """
+        Encoder followed by a projection MLP.
+        :param encoder: encoder to use
+        :param hidden_dim: hidden dimension
+        :param projection_out_dim: projection output dimension
+        """
+        super(EncoderProjecter, self).__init__()
+
+        self.encoder = encoder
+        self.projection = MLP(input_dim=128, projection_dim=projection_out_dim, hidden_dim=hidden_dim)
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.encoder(x)
+        return self.projection(h)
+
+
+# https://arxiv.org/pdf/2006.07733
+class BYOL(nn.Module):
+    def __init__(
+            self,
+            hidden_dim: int = 512,
+            projection_out_dim: int = 128,
+            target_decay: float = 0.9975
+    ):
+        """
+        BYOL model for self-supervised learning.
+        :param hidden_dim: hidden dimension
+        :param projection_out_dim: projection output dimension
+        :param target_decay: target network decay rate
+        """
+        super(BYOL, self).__init__()
+        encoder = VGGLikeEncode(in_channels=1, out_channels=128, feature_dim=32, apply_pooling=True)
+        self.online_network = EncoderProjecter(encoder)
+        self.online_predictor = MLP(input_dim=128, projection_dim=projection_out_dim, hidden_dim=hidden_dim)
+
+        self.target_network = EncoderProjecter(encoder)
+        self.target_network.load_state_dict(self.online_network.state_dict())
+
+        self.target_network.eval()
+        for param in self.target_network.parameters():
+            param.requires_grad = False
+        self.target_decay = target_decay
+        self.loss_function = NormalizedMSELoss()
+
+    @torch.no_grad()
+    def soft_update_target_network(self):
+        for online_p, target_p in zip(self.online_network.parameters(), self.target_network.parameters()):
+            target_p.data = target_p.data * self.target_decay + online_p.data * (1. - self.target_decay)
+
+    def forward(self, view: Tensor) -> Tuple[Tensor, Tensor]:
+        online_proj = self.online_network(view)
+        target_proj = self.target_network(view)
+
+        return online_proj, target_proj
+
+    def loss(self, view1: Tensor, view2: Tensor) -> Tensor:
+        online_proj1, target_proj1 = self(view1)
+        online_proj2, target_proj2 = self(view2)
+
+        online_prediction_1 = self.online_predictor(online_proj1)
+        online_prediction_2 = self.online_predictor(online_proj2)
+
+        loss1 = self.loss_function(online_prediction_1, target_proj2.detach())
+        loss2 = self.loss_function(online_prediction_2, target_proj1.detach())
+        return torch.mean(loss1 + loss2)

train_byol.py

@@ -0,0 +1,301 @@
+import argparse
+
+import torch
+import wandb
+from torch import nn, optim
+from torch.nn.functional import cosine_similarity
+from torch.optim import lr_scheduler
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from typing_extensions import Optional
+
+from src.dataset import RandomAugmentedDataset, get_byol_transforms
+from src.models import BYOL
+
+
+def get_data_loaders(
+        batch_size: int,
+        num_train_samples: int,
+        num_val_samples: int,
+        shape_params: Optional[dict] = None,
+        num_workers: int = 0
+):
+    augmentations = get_byol_transforms()
+
+    train_dataset = RandomAugmentedDataset(
+        augmentations,
+        shape_params,
+        num_samples=num_train_samples,
+        train=True
+    )
+    val_dataset = RandomAugmentedDataset(
+        augmentations,
+        shape_params,
+        num_samples=num_val_samples,
+        train=False
+    )
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers
+    )
+
+    return train_loader, val_loader
+
+
+def build_model(lr: float):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = BYOL().to(device)
+
+    optimizer = optim.Adam(
+        list(model.online_network.parameters()) + list(model.online_predictor.parameters()),
+        lr=lr
+    )
+    scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.1, patience=2)
+
+    return model, optimizer, scheduler, device
+
+
+def train_epoch(
+        model: nn.Module,
+        optimizer: optim.Optimizer,
+        train_loader: DataLoader,
+        device: torch.device
+) -> dict:
+    model.train()
+    running_train_loss = 0.0
+    total_cos_sim, total_l2_dist, total_feat_norm, total_grad_norm = 0.0, 0.0, 0.0, 0.0
+    num_train_batches = 0
+
+    for (view_1, view_2) in tqdm(train_loader, desc="Training"):
+        view_1 = view_1.to(device)
+        view_2 = view_2.to(device)
+
+        loss = model.loss(view_1, view_2)
+
+        optimizer.zero_grad()
+        loss.backward()
+
+        with torch.no_grad():
+            online_proj1, target_proj1 = model(view_1)
+            online_proj2, target_proj2 = model(view_2)
+
+            cos_sim = cosine_similarity(online_proj1, target_proj2).mean().item()
+            l2_dist = torch.norm(online_proj1 - target_proj2, dim=-1).mean().item()
+            feat_norm = torch.norm(online_proj1, dim=-1).mean().item()
+
+            grad_norm = torch.norm(
+                torch.cat([
+                    p.grad.flatten()
+                    for p in model.online_network.parameters()
+                    if p.grad is not None
+                ])
+            ).item()
+
+            total_cos_sim += cos_sim
+            total_l2_dist += l2_dist
+            total_feat_norm += feat_norm
+            total_grad_norm += grad_norm
+
+        optimizer.step()
+        model.soft_update_target_network()
+
+        running_train_loss += loss.item()
+        num_train_batches += 1
+
+    train_loss = running_train_loss / num_train_batches
+    train_cos_sim = total_cos_sim / num_train_batches
+    train_l2_dist = total_l2_dist / num_train_batches
+    train_feat_norm = total_feat_norm / num_train_batches
+    train_grad_norm = total_grad_norm / num_train_batches
+
+    return {
+        "loss": train_loss,
+        "cos_sim": train_cos_sim,
+        "l2_dist": train_l2_dist,
+        "feat_norm": train_feat_norm,
+        "grad_norm": train_grad_norm,
+    }
+
+
+@torch.no_grad()
+def validate(
+        model: nn.Module,
+        val_loader: DataLoader,
+        device: torch.device
+) -> dict:
+    model.eval()
+    running_val_loss = 0.0
+    total_cos_sim, total_l2_dist, total_feat_norm = 0.0, 0.0, 0.0
+    num_val_batches = 0
+
+    for (view_1, view_2) in tqdm(val_loader, desc="Validation"):
+        view_1 = view_1.to(device)
+        view_2 = view_2.to(device)
+
+        loss = model.loss(view_1, view_2)
+        running_val_loss += loss.item()
+
+        online_proj1, target_proj1 = model(view_1)
+        online_proj2, target_proj2 = model(view_2)
+
+        cos_sim = cosine_similarity(online_proj1, target_proj2).mean().item()
+        l2_dist = torch.norm(online_proj1 - target_proj2, dim=-1).mean().item()
+        feat_norm = torch.norm(online_proj1, dim=-1).mean().item()
+
+        total_cos_sim += cos_sim
+        total_l2_dist += l2_dist
+        total_feat_norm += feat_norm
+        num_val_batches += 1
+
+    val_loss = running_val_loss / num_val_batches
+    val_cos_sim = total_cos_sim / num_val_batches
+    val_l2_dist = total_l2_dist / num_val_batches
+    val_feat_norm = total_feat_norm / num_val_batches
+
+    return {
+        "loss": val_loss,
+        "cos_sim": val_cos_sim,
+        "l2_dist": val_l2_dist,
+        "feat_norm": val_feat_norm
+    }
+
+
+def train(
+        model: nn.Module,
+        optimizer: optim.Optimizer,
+        scheduler,
+        device: torch.device,
+        train_loader: DataLoader,
+        val_loader: DataLoader,
+        num_epochs: int,
+        early_stopping_patience: int = 3,
+        save_path: str = "best_byol.pth"
+):
+    best_loss = float("inf")
+    epochs_no_improve = 0
+    print("Start training...")
+
+    for epoch in range(num_epochs):
+        print(f"Epoch {epoch + 1}/{num_epochs}")
+
+        train_metrics = train_epoch(model, optimizer, train_loader, device)
+
+        val_metrics = validate(model, val_loader, device)
+
+        wandb.log({
+            "epoch": epoch + 1,
+            "train_loss": train_metrics["loss"],
+            "train_cos_sim": train_metrics["cos_sim"],
+            "train_l2_dist": train_metrics["l2_dist"],
+            "train_feat_norm": train_metrics["feat_norm"],
+            "train_grad_norm": train_metrics["grad_norm"],
+            "val_loss": val_metrics["loss"],
+            "val_cos_sim": val_metrics["cos_sim"],
+            "val_l2_dist": val_metrics["l2_dist"],
+            "val_feat_norm": val_metrics["feat_norm"],
+        })
+
+        print(
+            f"Train Loss: {train_metrics['loss']:.4f} | "
+            f"CosSim: {train_metrics['cos_sim']:.4f} | "
+            f"L2Dist: {train_metrics['l2_dist']:.4f}"
+        )
+        print(
+            f"Val Loss: {val_metrics['loss']:.4f} | "
+            f"CosSim: {val_metrics['cos_sim']:.4f} | "
+            f"L2Dist: {val_metrics['l2_dist']:.4f}"
+        )
+
+        current_val_loss = val_metrics["loss"]
+        if current_val_loss < best_loss or val_metrics['cos_sim'] >= 0.86:
+            best_loss = current_val_loss
+            encoder_state_dict = model.online_network.encoder.state_dict()
+            torch.save(encoder_state_dict, save_path)
+            epochs_no_improve = 0
+        else:
+            epochs_no_improve += 1
+
+        scheduler.step(val_metrics["cos_sim"])
+
+        if epochs_no_improve >= early_stopping_patience:
+            print(f"Early stopping on epoch {epoch + 1}")
+            break
+
+
+def main(config: dict):
+    wandb.init(project="contrastive_learning_byol", config=config)
+
+    train_loader, val_loader = get_data_loaders(
+        batch_size=config["batch_size"],
+        num_train_samples=config["num_train_samples"],
+        num_val_samples=config["num_val_samples"],
+        shape_params=config["shape_params"]
+    )
+
+    model, optimizer, scheduler, device = build_model(
+        lr=config["lr"]
+    )
+
+    train(
+        model=model,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        device=device,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        num_epochs=config["num_epochs"],
+        early_stopping_patience=config["early_stopping_patience"],
+        save_path=config["save_path"]
+    )
+
+    wandb.finish()
+
+
+if __name__ == "__main__":
+    # parser = argparse.ArgumentParser(description="Train BYOL model")
+    # parser.add_argument("--batch_size", type=int, default=512)
+    # parser.add_argument("--lr", type=float, default=5e-4)
+    # parser.add_argument("--num_epochs", type=int, default=15)
+    # parser.add_argument("--num_train_samples", type=int, default=100000)
+    # parser.add_argument("--num_val_samples", type=int, default=10000)
+    # parser.add_argument("--random_intensity", type=int, default=1)
+    # parser.add_argument("--early_stopping_patience", type=int, default=3)
+    # parser.add_argument("--save_path", type=str, default="best_byol.pth")
+    # args = parser.parse_args()
+
+    # config = {
+    #     "batch_size": args.batch_size,
+    #     "lr": args.lr,
+    #     "num_epochs": args.num_epochs,
+    #     "num_train_samples": args.num_train_samples,
+    #     "num_val_samples": args.num_val_samples,
+    #     "shape_params": {
+    #         "random_intensity": bool(args.random_intensity)
+    #     },
+    #     "early_stopping_patience": args.early_stopping_patience,
+    #     "save_path": args.save_path
+    # }
+
+    config = {
+        "batch_size": 1024,
+        "lr": 1e-3,
+        "num_epochs": 15,
+        "num_train_samples": 100000,
+        "num_val_samples": 10000,
+        "shape_params": {
+            "random_intensity": True
+        },
+        "early_stopping_patience": 3,
+        "save_path": "best_byol.pth"
+    }
+
+    main(config)

train_cross_classifier.py

@@ -0,0 +1,295 @@
+import argparse
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+from torch.optim.lr_scheduler import StepLR
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from typing_extensions import Optional
+
+from src.dataset import RandomPairDataset
+from src.models import CrossAttentionClassifier, VGGLikeEncode
+
+
+def visualize_attention(attn_heatmap, epoch: int):
+    fig, ax = plt.subplots(figsize=(6, 6))
+    im = ax.imshow(attn_heatmap, cmap="hot", interpolation="nearest")
+    plt.colorbar(im, fraction=0.046, pad=0.04)
+    plt.title(f"Attention Heatmap (Flatten 64x64) | Epoch {epoch}")
+
+    wandb.log({"Flatten Attention Heatmap": wandb.Image(fig, caption=f"Flatten 64x64 | Epoch {epoch}")})
+
+    plt.close(fig)
+
+
+def get_data_loaders(
+        num_train_samples: int,
+        num_val_samples: int,
+        batch_size: int,
+        num_workers: int = 0,
+        shape_params: Optional[dict] = None,
+):
+    train_dataset = RandomPairDataset(
+        shape_params=shape_params,
+        num_samples=num_train_samples,
+        train=True
+    )
+    val_dataset = RandomPairDataset(
+        shape_params=shape_params,
+        num_samples=num_val_samples,
+        train=False
+    )
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers
+    )
+
+    return train_loader, val_loader
+
+
+def build_model(
+        path_to_encoder: str,
+        lr: float,
+        weight_decay: float,
+        step_size: int,
+        gamma: float,
+        device: torch.device
+):
+    encoder = VGGLikeEncode(in_channels=1, out_channels=128, feature_dim=32, apply_pooling=False)
+    encoder.load_state_dict(torch.load(path_to_encoder))
+
+    model = CrossAttentionClassifier(encoder=encoder)
+    model = model.to(device)
+
+    criterion = nn.BCEWithLogitsLoss()
+
+    optimizer = optim.Adam(
+        model.parameters(),
+        lr=lr,
+        weight_decay=weight_decay
+    )
+
+    scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma)
+
+    return model, criterion, optimizer, scheduler
+
+
+def train_epoch(
+        model: nn.Module,
+        criterion: nn.Module,
+        optimizer: optim.Optimizer,
+        train_loader: DataLoader,
+        device: torch.device
+):
+    model.train()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+
+    for img1, img2, labels in tqdm(train_loader, desc="Training", leave=False):
+        img1, img2, labels = img1.to(device), img2.to(device), labels.to(device)
+
+        optimizer.zero_grad()
+
+        logits, attn_weights = model(img1, img2)
+        loss = criterion(logits, labels)
+
+        loss.backward()
+        optimizer.step()
+
+        running_loss += loss.item() * img1.size(0)
+
+        preds = (torch.sigmoid(logits) > 0.5).float()
+        correct += (preds == labels).sum().item()
+        total += labels.size(0)
+
+    epoch_loss = running_loss / len(train_loader.dataset)
+    epoch_acc = correct / total
+
+    return epoch_loss, epoch_acc
+
+
+@torch.no_grad()
+def validate(
+        model: nn.Module,
+        criterion: nn.Module,
+        val_loader: DataLoader,
+        device: torch.device
+):
+    model.eval()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+
+    for img1, img2, labels in tqdm(val_loader, desc="Validation", leave=False):
+        img1, img2, labels = img1.to(device), img2.to(device), labels.to(device)
+
+        logits, attn_weights = model(img1, img2)
+        loss = criterion(logits, labels)
+
+        running_loss += loss.item() * img1.size(0)
+
+        preds = (torch.sigmoid(logits) > 0.5).float()
+        correct += (preds == labels).sum().item()
+        total += labels.size(0)
+
+    epoch_loss = running_loss / len(val_loader.dataset)
+    epoch_acc = correct / total
+
+    return epoch_loss, epoch_acc
+
+
+def train(
+        model: nn.Module,
+        criterion: nn.Module,
+        optimizer: optim.Optimizer,
+        scheduler,
+        train_loader: DataLoader,
+        val_loader: DataLoader,
+        device: torch.device,
+        num_epochs: int = 30,
+        save_path: str = "best_attention_classifier.pth"
+):
+    best_val_loss = float("inf")
+    epochs_no_improve = 0
+    print("Start training...")
+
+    for epoch in range(num_epochs):
+        print(f"Epoch {epoch + 1}/{num_epochs}")
+
+        train_loss, train_acc = train_epoch(model, criterion, optimizer, train_loader, device)
+
+        val_loss, val_acc = validate(model, criterion, val_loader, device)
+
+        scheduler.step()
+
+        wandb.log({
+            "epoch": epoch + 1,
+            "train_loss": train_loss,
+            "train_acc": train_acc,
+            "val_loss": val_loss,
+            "val_acc": val_acc,
+            "lr": optimizer.param_groups[0]["lr"],
+        })
+
+        print(
+            f"learning rate: {optimizer.param_groups[0]['lr']:.6f}, "
+            f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, "
+            f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}"
+        )
+
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            torch.save(model.state_dict(), save_path)
+            epochs_no_improve = 0
+        else:
+            epochs_no_improve += 1
+
+        with torch.no_grad():
+            sample_img1, sample_img2, sample_labels = next(iter(val_loader))
+            sample_img1, sample_img2 = sample_img1.to(device), sample_img2.to(device)
+
+            _, sample_attn_weights = model(sample_img1, sample_img2)
+
+            wandb.log({
+                "attention_std": sample_attn_weights.std().item(),
+                "attention_mean": sample_attn_weights.mean().item(),
+            })
+
+            attn_heatmap = sample_attn_weights[0].detach().cpu().numpy()
+            visualize_attention(attn_heatmap, epoch)
+
+
+def main(config):
+    wandb.init(project="cross_attention_classifier", config=config)
+
+    train_loader, val_loader = get_data_loaders(
+        shape_params=config["shape_params"],
+        num_train_samples=config["num_train_samples"],
+        num_val_samples=config["num_val_samples"],
+        batch_size=config["batch_size"]
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model, criterion, optimizer, scheduler = build_model(
+        path_to_encoder=config["path_to_encoder"],
+        lr=config["lr"],
+        weight_decay=config["weight_decay"],
+        step_size=config["step_size"],
+        gamma=config["gamma"],
+        device=device
+    )
+
+    train(
+        model=model,
+        criterion=criterion,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        device=device,
+        num_epochs=config["num_epochs"],
+        save_path=config["save_path"]
+    )
+
+    wandb.finish()
+
+
+if __name__ == "__main__":
+
+    # parser = argparse.ArgumentParser(description="Train classifier model")
+    # parser.add_argument("--path_to_encoder", type=str, default="best_byol.pth")
+    # parser.add_argument("--batch_size", type=int, default=256)
+    # parser.add_argument("--lr", type=float, default=8e-5)
+    # parser.add_argument("--weight_decay", type=float, default=1e-4)
+    # parser.add_argument("--step_size", type=int, default=10)
+    # parser.add_argument("--gamma", type=float, default=0.1)
+    # parser.add_argument("--num_epochs", type=int, default=10)
+    # parser.add_argument("--num_train_samples", type=int, default=10000)
+    # parser.add_argument("--num_val_samples", type=int, default=2000)
+    # parser.add_argument("--save_path", type=str, default="best_attention_classifier.pth")
+    # args = parser.parse_args()
+
+    # config = {
+    #     "path_to_encoder": args.path_to_encoder,
+    #     "batch_size": args.batch_size,
+    #     "lr": args.lr,
+    #     "weight_decay": args.weight_decay,
+    #     "step_size": args.step_size,
+    #     "gamma": args.gamma,
+    #     "num_epochs": args.num_epochs,
+    #     "num_train_samples": args.num_train_samples,
+    #     "num_val_samples": args.num_val_samples,
+    #     "save_path": args.save_path,
+    # }
+
+    config = {
+        "path_to_encoder": "best_byol.pth",
+        "batch_size": 256,
+        "lr": 8e-5,
+        "weight_decay": 1e-4,
+        "step_size": 10,
+        "gamma": 0.1,
+        "num_epochs": 10,
+        "num_train_samples": 10000,
+        "num_val_samples": 2000,
+        "save_path": "best_attention_classifier.pth",
+    }
+
+    if "shape_params" not in config:
+        config["shape_params"] = {}
+
+    main(config)