Код на JAX/Flax и предобученные чекпойнты моделей Vision Transformer (ViT) и MLP-Mixer для классификации изображений из статей Google Research.
В репозитории опубликованы код и предобученные чекпойнты Vision Transformer (ViT) и MLP-Mixer — двух архитектур для распознавания изображений, которые обрабатывают изображения как последовательности патчей фиксированного размера или как сети, состоящие только из MLP. Репозиторий сопровождает несколько статей Google Research, включая оригинальную статью о ViT, работу MLP-Mixer, исследование AugReg по обучению ViT, а также более поздние работы о LiT — моделях для zero-shot сопоставления изображений и текста — и обучении с учётом резкости.
Модели предобучены на ImageNet и ImageNet-21k, а код включает скрипты для дообучения на JAX и Flax на таких наборах данных, как CIFAR-10/100, и на любых наборах, доступных через TensorFlow Datasets. Многие чекпойнты размещены в бакетах Google Cloud Storage; среди них — более 50 тысяч вариантов ViT и гибридных моделей с разными настройками аугментации и регуляризации. В Colab-ноутбуках рассматриваются модели, их дообучение и инференс LiT; несколько чекпойнтов также можно загрузить с помощью библиотеки timm для PyTorch.
Проект не является официальным продуктом Google. По словам автора, более продвинутый код для обучения, например для обучения на нескольких хостах, находится в отдельном репозитории big_vision.
Убедитесь, что установлена Python версии 3.10 или новее
Установите JAX и зависимости для GPU
pip install -r vit_jax/requirements.txtДообучите модель ViT-B/16 (предварительно обученную на ImageNet-21k) на CIFAR-10
python -m vit_jax.main --workdir=/tmp/vit-$(date +%s) --config=$(pwd)/vit_jax/configs/vit.py:b16,cifar10 --config.pretrained_dir='gs://vit_models/imagenet21k'Описание создано ИИ на основе README проекта. Статистика GitHub обновляется ежедневно.
Библиотека PyTorch с архитектурами моделей для обработки изображений, предобученными весами и скриптами обучения, валидации и экспорта для задач компьютерного зрения.
Реализация Vision Transformer (ViT) на PyTorch и множества его вариантов для классификации изображений и других задач компьютерного зрения.
Библиотека предобученных моделей для Keras 3: модели для текста, изображений и аудио, работающие с JAX, TensorFlow или PyTorch.
Утилиты для сохранения и восстановления состояния моделей JAX, например весов и состояния оптимизатора.
Модульная масштабируемая библиотека для обучения моделей машинного обучения на JAX/Flax, созданная для быстрых исследовательских итераций.
Ultralytics YOLO — библиотека Python для компьютерного зрения: обнаружение объектов, сегментация, классификация, оценка позы, трекинг и другие задачи.