PyTorch Hub: база предобученных моделей на PyTorch

PyTorch Hub — это открытая база с предобученными нейросетевыми моделями на PyTorch. Сейчас в базе 19 моделей, в том числе для задач компьютерного зрения и обработки естественного языка. PyTorch Hub оформлен в формате API. С помощью API исследователи публикуют модели, а пользователи используют ранее опубликованные модели.

Воспроизводимость исследований в машинном обучении — актуальная проблема. Большинство публикаций либо сложно воспроизвести, либо невозможно из-за недостаточной детализации условий экспериментов и отсутствия кода. PyTorch Hub решает проблему для исследователей, которые хотят опубликовать свою модель, и для пользователей, которые хотят воспользоваться какой-то моделью для решения своих задач.

PyTorch Hub поддерживает Google Colab и интеграцию с сервисом Papers With Code. Сейчас среди опубликованных нейросетей модели для классификации, сегментации, генеративные и трансформеры.

Публикация модели

PyTorch Hub поддерживает публикацию предобученных моделей (код модели + веса) в репозиторий на GitHub. Для этого нужно добавить в репозиторий один .py файл. В этом файле хранится список библиотек, которые нужны для обучения модели. Примеры таких файлов доступны для опубликованных в Hub-е моделей: torchvision, huggingface-bert и gan-model-zoo.

Как только файл создан, нужно отправить pull-запрос в официальный репозиторий PyTorch Hub (пример). Команда PyTorch фильтрует модели, которые им предлагают для публикации, поэтому затем идет скрининг. После этого модель появляется на официальной странице базы.

Использование предобученной модели

Пользователю доступен следующий функционал:

• получить актуальный список опубликованных моделей;
• загрузить модель;
• получить список доступных методов для загруженной модели

У опубликованной модели PyTorch Hub могут быть методы для предобработки. Например, у BERT есть метод для токенизации текстов.

Модели, доступные в PyTorch Hub, напрямую связаны с Papers With Code и поддерживают Google Colab.