В этой домашке вам предстоит сделать следующие вещи:

1. Реализовать свой маленький ViT;
2. Обучить его на задаче классификации головных уборов. Не гонитесь за качеством на валидации или на тесте, просто убедитесь, что модель учится на тренировочном наборе данных (лосс падает);
3. [*] Реализовать трюк с ускорением подсчета self-attention;
4. [*] Заставить обученный ViT работать на другом разрешении.

Ваша задача написать код модели ViT, основываясь на

1. ViT paper;
2. Статья от DeepSchool. В целом, вы можете пользоваться любыми материалами, лишь бы вам было понятно то, что вы делаете, и в конце концов у вас собралась модель.

Мы ожидаем, что вы реализуете модель модульно, т.е. у вас будут следующие блоки:

1. PatchEmbedder - модуль, отвечающий за перевод изображения в патчи. Можете сделать хоть сверточный вариант, хоть линейный. Делайте, какой больше нравится;
2. LinearProjector - на вход изображение, а на выходе патчи в виде векторов, сложенные с позиционными эмбеддингами;
3. ScaledDotProductAttention - модуль, считающий атеншн; 4 MultiHeadSelfAttention - модуль, содержащий в себе несколько голов ScaledDotProductAttention;
4. EncoderBlock - модуль, реализующий один блок энкодера ViT.
5. ViT - модуль, реализующий сам ViT с помощью всех предыдущих блоков.

Конечно, это необязательно требование, и вы можете реализовать модельку как вашей душе угодно, но проверяющим будет проще, если у вас будут эти блоки.

Вам нужно сделать следующее:

1. Взять вот этот набор данных головных уборов. Он очень маленький, скачать его не составит труда;
2. Взять вашу реализацию модели ViT из предыдущего пунка;
3. Написать обучалку, как вас учили на первых лекциях с блекджеком, логгингом в ClearML и прочими простыми радостями мл инженера;
4. Запустить обучалку и убедиться, что модель учится. Не нужно пытаться выбивать заоблачную метрику на валидационной\тестовой выборках. Нужно просто убедиться, что моделька учится и получается какой-то accuracy > 0.05. В качестве гиперпараметров для обучения можно взять настройки прям из исходной статьи по ViT.

Мы считаем Q, K, V отдельными линейными слоями с весами W_q, W_k, W_v. Можно сделать умнее и посчитать их одним линейным слоем с расширенной матрицей весов W. Схематически это отображено на следующей картинке:

Представим, что вы обучили ваш супер-дупер-фенси-шменси ViT на разрешении 224x224x3. К вам приходит ваш бородатый тимлид и говорит, что инференс хотим делать на изображениях размера 384x384x3, но вот незадача - ваш ViT из коробки не может работать на изображениях другого разрешения. Ваша задача внимательно прочитать статьи, перечисленные в задании по реализации ViT, понять, как эту проблему предлагают решать авторы оригинальной статьи и реализовать этот способ.