Совсем недавно компания Сбер представила ряд, адаптированных под русский язык, моделей машинного обучения под названием RuCLIP. Суть их работы заключается в возможности сравнения схожести текстового описания и изображения. Рассматривая результаты тестирования, мы заметили, что их модели довольно большие (150+ миллионов параметров) и при этом занимают довольно много места, связи с чем решили сделать свою маленькую, быструю и масштабируемую версию под названием RuCLIP tiny.
Данные
В наших экспериментах мы попробовали датасеты разного качества и остановились на выборе переведённых наборов COCO и CC12M. Они оба содержат пары текст + картинка. На тот момент CC12M ещё не был переведён, связи с чем пришлось обработать его с использованием гугл переводчика. датасет cc12m гугл переводчиком.
Выбор архитектуры
Архитектура CLIP’а состоит из энкодера для картинок и энкодера для текста. Их задача состоит в том, чтобы уменьшить косинусное расстояние между их выходными векторами.
В начале мы сделали архитектуру с cointegrated/rubert-tiny в качестве текстового энкодера и swin_tiny_patch4_window7_224 в качестве энкодера изображений. Но она показала себя не очень относительно следующей архитектуры.
Тогда мы решили обновить архитектуру. Для текстового энкодера мы выбрали недавно вышедший DeepPavlov/distilrubert-tiny-cased-conversational-v1, а для энкодера изображений - convnext_tiny, который быстрее, чем swin transformer.
Всего в новой архитектуре 38317537 параметров.
Тренировка модели
Для тренировки мы заморозили энкодер изображений и обучали только энкодер текста, как это делали в этой статье. Там они показали, что когда энкодер изображений был предварительно обучен, то необязательно производить его файнтюинг во время обучения всей модели.
Поскольку мы тренировали на гугл колаб про, то у нас была всего одна NVIDIA Tesla P100 на 16 гб. К сожалению, у нас не получалось сделать размер батча выше, чем 256, а для contrastive loss он очень важен. Тогда мы стали искать различные статьи и наткнулись на эту статью. Там авторы рассказывают, как реализовать накопление градиента для contrastive loss. Мы это реализовали и смогли поставить размер батча равный 2048.
Результаты модели
Для оценки модели мы решили взять Google Colab Pro + с NVIDIA Tesla V100.
На датасете CIFAR100 мы получили 0.4662 top1 accuracy и 0.7318 top5 accuracy.
Скорость модели
batch size | encode_image | encode_text | total |
2 | 11мс | 4мс | 15мс |
8 | 12мс | 4мс | 16мс |
16 | 13мс | 3мс | 16мс |
32 | 12мс | 4мс | 16мс |
64 | 13мс | 4мс | 17мс |
Подведение итогов
В результате нашей работы мы сделали ru-CLIP-tiny, который способен оценивать схожесть картинок с текстом. Код и веса вы можете найти тут.
Хочется выразить благодарность Sber AI за предоставленные гранты, на которые были проведены исследования, в рамках Artificial Intelligence International Junior Contest (AIIJC).
