Пример создания сервиса классификации изображений

В данном репозитории реализован код и инструкции для создания и запуска сервиса классификации изображений: от этапа обучения модели до деплоя в виде API сервиса в Docker-контейнере.

Обучение модели

Перед запуском процесса обучения необходимо установить зависимости:

pip install -r dependencies/train.requirements.txt

Код для загрузки датасета, создания и обучения модели и её конвертации в ONNX приведён в jupyter-ноутбуке notebooks/train.ipynb.

В результате выполнения всех ячеек в ноутбуке в папке с проектом должны появиться директории dataset и model. В директории dataset будут располагаться подпапки с датасетами для обучения, валидации и тестирования модели. В директории model будут располагаться файлы model_state_dict.pth (файл с весами обученной модели) и model.onnx (модель в формате ONNX).

Настройка DVC

Перед настройкой DVC необходимо создать внешнее хранилище данных и убедиться, что в вашем проекте инициализирован git. Для данного примера воспользуемся Google Drive. Для настройки DVC с хранилищем в Google Drive необходимо выполнить следующее:

1. Выполнить команду инициализации DVC:

dvc init

2. Создать папку в Google Drive.
3. Зайти в созданную папку и скопирать идентификатор папки из URL (в адресной строке браузера).

Например, для URL https://drive.google.com/drive/folders/0AIac4JZqHhKmUk9PDA идентификатором будет являться 0AIac4JZqHhKmUk9PDA.

4. Добавить папку в Google Drive в качестве удалённого хранилища для DVC:

dvc remote add --default origin gdrive://0AIac4JZqHhKmUk9PDA
dvc remote modify origin gdrive_acknowledge_abuse true

Загрузка модели в удалённое хранилище через DVC

Для того, чтобы загрузить модель в Google Drive через DVC, воспользуйтесь следующими командами:

dvc add model
dvc push

Во время выполнения команды dvc push будет необходимо пройти аутентификации в Google (автоматически откроется вкладка в браузере). В папке с проектом должен появиться файл model.dvc и в файл .gitignore должна записаться строчка /model.

В созданной в Google Drive папке должна появиться директория files/md5.

Скачивание модели из удалённого хранилища через DVC

Для загрузки модели при наличии в репозитории папки с расширением .dvc и файла model.dvc можно воспользоваться следующей командой:

dvc pull

Сборка Docker-образа

Можно собрать два варианта Docker-образов: для инференса модели на CPU или на GPU.

Команда для сборки образа под запуск на CPU:

docker build -f cpu.Dockerfile -t cat-dog-classification-service:cpu .

Команда для сборки образа под запуск на GPU:

docker build -f gpu.Dockerfile -t cat-dog-classification-service:gpu .

Внимание! Со сборкой контейнера для инференса на GPU есть некоторые особенности. В частности, требуется знать версию CUDA, которая установлена на вашем ПК.

Проверить версию CUDA можно выполнив следующую команду:

nvidia-smi

Узнав версию CUDA, проверьте, поддерживает ли её выбранная версия ONNX Runtime (по умолчанию в репозитории приведена 1.17) на сайте.

Если у вас версия CUDA 12.2 и старше, то gpu.Dockerfile можно оставить без изменений.

Если у вас вресия CUDA 11.8 и младше, то подберите подходящую версию библиотеки ONNX Runtime на сайте и замените её в gpu.Dockerfile (строка 14). Также найди подходящий Docker-образ от Nvidia на сайте, введя в поле с фильтром вашу версию CUDA. Выбрать надо ту версию, в названии которой есть слова runtime и ubuntu22.04.

Запуск Docker-контейнера с сервисом

Запуск Docker-контейнера может быть из двух вариантов Docker-образа: для инференса модели на CPU или на GPU.

Команда для запуска на CPU:

docker run -d --name cat-dog-classification-service \
-e SERVICE_NAME="Cat-Dog classification service" \
-p 8080:8000 \
cat-dog-classification-service:cpu

Команда для запуска на GPU:

docker run -d --name cat-dog-classification-service \
--gpus all \
-e SERVICE_NAME="Cat-Dog classification service" \
-p 8080:8000 \
cat-dog-classification-service:gpu

Переменную окружения SERVICE_NAME можно передать опционально. После запуска контейнера можно перейти по ссылке http://SERVER_IP_OR_HOSTNAME:8080/docs и поизучать OpenAPI-документацию.

P.S. SERVER_IP_OR_HOSTNAME - ip-адрес или имя вашего сервера (для запуска локально - locahost).

Остановка Docker-контейнера с сервисом

Тут одна команда для любого типа инференса:

docker stop cat-dog-classification-service

Удаление Docker-контейнера с сервисом

Тут одна команда для любого типа инференса:

docker rm cat-dog-classification-service