Skip to content

Code for training deep learning models for dog attributes recognition.

Notifications You must be signed in to change notification settings

KEKDATA/KEKPETS_ATTRS

Repository files navigation

О задаче

Код в этом репозитории предназначен для тренировки 2 моделей для решения следующих задач:

  1. Задача классификации длины хвоста собаки (длинный, короткий);
  2. Задача классификации оттенка шерсти собаки (темный, светлый, рыжий, мультицвет).

Ниже в отдельных секциях будут даны инструкции по каждой из задач в отдельности.

О данных

Данные содержат в себе информацию об оттенке шерсти собак и о длине их хвостов.

Формирование набора данных

Набор данных для этой модели был подготовлен следующим образом:

  1. Из исходных изображений были вырезаны собаки;
  2. Кропы, полученные на предыдущем шаге, были вручную классифицированы на категории по цвету и по длине хвоста.

Цвет шерсти

Изначально было дано три оттенка цвета шерсти: темный, светлый и мультицвет. Мы выделили дополнительно рыжий оттенок в отдельный класс, так как разметка была очень плохой и рыжий оттенок находился во всех трех изначальных классах. Итоговое множество цветов животного:

  1. dark - темный окрас шерсти;
  2. bright - светлый окрас шерсти;
  3. ginger - рыжий окрас шерсти;
  4. multicolor - окрас шерсти содержит несколько цветов.

Длины хвостов

Данные, полученные от организаторов, содержат огромное число ошибок. Поэтому мы переразметили кропы вручную.

Организация репозитория

  • Ноутбук create_csv.ipynb содержит в себе код подготовки .csv файла для процесса тренировки модели;
  • Модуль dataset/animal.py содержит в себе код для работы с Pytorch-датасетом;
  • Модуль model/animal.py содержит в себе код модели для решения поставленной задачи;
  • Модуль utils/augmentation.py содержит в себе вспомогательные функции для аугментации данных;
  • Модуль utils/config.py содержит в себе вспомогательные функции для конфигурации проекта;
  • Модуль utils/data.py содержит в себе вспомогательные функции для работы с данными;
  • Модуль utils/ema.py содержит в себе код экспоненциально скользящего среднего;
  • Модуль utils/general.py содержит в себе общие вспомогательные функции;
  • Модуль utils/training.py содержит в себе функции процесса тренировки;
  • Модуль utils/validation.py содержит в себе функции оценки качества модели;
  • Модуль train.py является точкой запуска процесса тренировки модели.
  • Файл config.yml является конфигурационном файлом, в котором находятся все настройки проекта.
  • Файл export_to_onnx.py содержит в себе код перевода модели в формат ONNX и оптимизацию производительности модели.

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

Требования для запуска проекта

Проект рассчитан на работу с python3.8. Зависимости, необходимые для запуска проекта, указаны в файле requirements.txt.

Как подготовить свой набор данных для обучения?

Вам необходимо собрать фотографии собак, на которых будет виден цвет шерсти собаки и ее хвост. После того, как фотографии собак будут собраны, вам необходимо создать .csv файл с путями до этих фотографий и метками классов. Пример такого файла вы можете найти в example_csv/train.csv.

Запуск проекта

Общие настройки

  1. Склонируйте себе этот репозиторий и перейдите в него.
  2. Создайте свежее виртуальное окружение и активируйте его: инструкция;
  3. Установите необходимые зависимости: pip install -r requirements.txt;
  4. Воспользуйтесь ClearML;
  5. Измените файл config.yml, выставив желаемые вами настройки;
  6. Запустите файл train.py.

Файл с весами лучшей модели появится в директории outputs/. После получения файла модели переведите ее в формат ONNX следующей командой: python export_to_onnx.py -i %path_to_weights% -o %path_to_onnx_model% -s 224,224 Переведенная в ONNX модель будет подвержена очень сильной оптимизации графа вычислений, поэтому скорость ее работы очень сильно увеличится. Замечание: Модель после процедуры оптимизации начинает работать очень быстро без потери качества, поэтому рекомендуем воспользоваться процедурой оптимизации модели. Если вы не хотите самостоятельно править конфигурационные файлы, то в папке predefined_configs содержатся конфигурационный файлы для каждой из задач классификации.

Результаты

Ниже дан краткий обзор получившихся результатов раздельно по каждой из задач.

Задача классификации длины хвоста

Ссылка на лучший эксперимент. По этой ссылке можно найти абсолютно всю информацию о процессе обучения, включая все гиперпараметры.

Для оценки качества модели использовалась метрика ROC AUC. Была взята именно эта метрика, т.к. она является метрикой ранжирования, т. е. напрямую смотрит на вероятности. Нам важны именно вероятности, потому что организаторы хотят проверить гипотезу о поиске животного по признакам. Признак, представляющий собой вероятность, а не хард-метку, гораздо лучше подходит для поисковых\матчинговых задач, отражает некоторую степень принадлежности, а не категоричное бинарное состояние. Лучшая модель дала на тестовом множестве ROC AUC = 0.88.

Задача классификации оттенка шерсти

Ссылка на лучший эксперимент. По этой ссылке можно найти абсолютно всю информацию о процессе обучения, включая все гиперпараметры.

Для оценки качества модели использовалась метрика F1. Лучшая модель дала на тестовом множестве F1 = 0.85.

Конфигурация машины, на которой производилась тренировка

  1. CPU: 6 Core Intel CPU;
  2. GPU: NVidia Tesla K80 11GB;
  3. RAM: 56 GB;
  4. OS: Ubuntu Linux 20.04. Время тренировки моделей в течение 15 эпох на переразмеченном нами датасете занимает 07:49 и 05:11 минут для задач классификации оттенка шерсти и классификации длины хвоста соответственно. Пайплайн построен таким образом, чтобы максимально удобным и быстрым образом производить эксперименты.

Конфигурация для инференса.

Модели могут использоваться даже на слабых процессорах в ONNX формате. На машине с 1 ядром и 2000 МБ оперативной памяти каждая из моделей инферится примерно за 170 милисекунд. Такая скорость достигается засчет легкой архитектуры и оптимизации графа вычислений.

Где найти натренированные модели?

  1. Модель классификации хвоста доступна по этой ссылке;
  2. Модель классификации цвета доступна по этой ссылке.

About

Code for training deep learning models for dog attributes recognition.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published