Код в этом репозитории предназначен для тренировки 2 моделей для решения следующих задач:
- Задача классификации длины хвоста собаки (длинный, короткий);
- Задача классификации оттенка шерсти собаки (темный, светлый, рыжий, мультицвет).
Ниже в отдельных секциях будут даны инструкции по каждой из задач в отдельности.
Данные содержат в себе информацию об оттенке шерсти собак и о длине их хвостов.
Набор данных для этой модели был подготовлен следующим образом:
- Из исходных изображений были вырезаны собаки;
- Кропы, полученные на предыдущем шаге, были вручную классифицированы на категории по цвету и по длине хвоста.
Изначально было дано три оттенка цвета шерсти: темный, светлый и мультицвет. Мы выделили дополнительно рыжий оттенок в отдельный класс, так как разметка была очень плохой и рыжий оттенок находился во всех трех изначальных классах. Итоговое множество цветов животного:
dark
- темный окрас шерсти;bright
- светлый окрас шерсти;ginger
- рыжий окрас шерсти;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
.
- Склонируйте себе этот репозиторий и перейдите в него.
- Создайте свежее виртуальное окружение и активируйте его: инструкция;
- Установите необходимые зависимости:
pip install -r requirements.txt
; - Воспользуйтесь ClearML;
- Измените файл
config.yml
, выставив желаемые вами настройки; - Запустите файл
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.
- CPU: 6 Core Intel CPU;
- GPU: NVidia Tesla K80 11GB;
- RAM: 56 GB;
- OS: Ubuntu Linux 20.04. Время тренировки моделей в течение 15 эпох на переразмеченном нами датасете занимает 07:49 и 05:11 минут для задач классификации оттенка шерсти и классификации длины хвоста соответственно. Пайплайн построен таким образом, чтобы максимально удобным и быстрым образом производить эксперименты.
Модели могут использоваться даже на слабых процессорах в ONNX
формате. На машине с 1 ядром и 2000 МБ оперативной памяти каждая из моделей инферится примерно за 170 милисекунд. Такая скорость достигается засчет легкой архитектуры и оптимизации графа вычислений.