Классификация изображений свёрточной нейросетью
В этом проекте мы обучаем свёрточную нейросеть классифицировать изображения нескольких классов. На выходе мы получаем обученный алгоритм, который принимает на вход изображение и возвращает прогноз класса, к которому оно может принадлежать.
Подготовим окружение
Установим необходимые пакеты и настроим среду на вывод статичных графиков. В случае проблем с библиотеками выполните команду EngeePkg.purge(), которая устранит из вашей системы все пакеты кроме системных и позволит установить необходимые пакеты без проблем с совместимостью.
# Установка необходимых пакетов
Pkg.add(["Flux", "BSON", "ImageTransformations"])
gr();
Обучающие данные расположены в нескольких папках в каталоге "учебные данные". Названия папок считываются из файловой системы и становятся названиями классов. Мы также указываем, где находятся примеры для классификации (с неизвестными метками).
DATA_DIR = "$(@__DIR__)/учебные данные";
UNKNOWN_DIR = "$(@__DIR__)/неизвестно";
Обучение и проверка модели
В этом проекте с нуля мы обучим сверточную модель следующего вида:
.svg)
В скрипте train.jl собраны функции, которые осуществляют:
- Загрузку изображений из папок классов, ресайз до 128×128, нормализация
- Подсчёт дисбаланса классов (сколько вилок, сколько ложек)
- Разделение данных на train/test (стратифицированно, 75%/25%)
- Сборка модели — свёрточная сеть с BatchNorm и Dropout
- Сбалансированные батчи — чтобы в каждом батче были оба класса поровну
- Аугментация — удвоение образцов в датасете при помощи отражения и высветления
- Цикл обучения — forward pass, подсчёт потерь, backward pass, обновление весов
- Оценка качества — точность, precision, recall на тесте
- Early stopping — остановка если нет улучшений 8 эпох
- Сохранение лучшей модели
include("$(@__DIR__)/_scripts/train.jl")
model, classes = train_model(DATA_DIR; epochs=25, batch_size=16, lr=0.0005, test_split=0.25);
После обучения мы преврим, как ведет себя модель. Скрипт выполняет следующие шаги:
-
Загрузка модели, метаданных и неизвестных изображений — из BSON-файла извлекается архитектура сети, веса, названия классов и количество примеров в каждом классе обучающего набора, и переключение модели в режим оценки. Сканируется указанная папка, каждое изображение приводится к размеру 128×128, нормализуется в диапазон [-1, 1] и переводится в формат тензора W×H×C
-
Классификация каждого изображения — тензор подаётся на вход модели, выходные логиты преобразуются в вероятности через softmax, определяется класс с максимальной вероятностью и соответствующая уверенность и группировка по классам
-
Сортировка по уверенности и вывод статистики — внутри каждого класса изображения сортируются от самых уверенных предсказаний к наименее уверенным. Количество предсказанных изображений, процент от общего числа, средняя, максимальная и минимальная уверенность
-
Анализ смещения — сравнение процента предсказаний каждого класса с процентом этого класса в обучающем наборе, вывод индикатора
include("$(@__DIR__)/_scripts/visualize.jl")
results = classify_and_visualize(UNKNOWN_DIR);
Классифицируем все изображения из папки "неизвестно" и выведем до 10 изображений для каждого класса.
include("$(@__DIR__)/_scripts/simple_mosaic.jl")
plot(create_simple_mosaic(UNKNOWN_DIR))
.png)
А следующий скрипт позволяет заново открыть обученную модель, прогнозировать класс для всех файлов из каталога неразмеченных примеров и вывести результаты в таблицу.
include("$(@__DIR__)/_scripts/predict_to_csv.jl")
predict_to_csv(UNKNOWN_DIR, confidence_threshold=0.6, output_csv="$(@__DIR__)/predictions.csv")
Заключение
В проекте построен фундамент для простой задачи: расположить изображения в нескольких папках и обучить классификатор, который позволит распределить по тем же классам изображения из второй папки с неизвестными объектами.
Результаты, представленные в примере, получены при помощи нейросети, прошедшей десятки итераций процесса обучения. Чтобы получить хороший алгоритм классификации, нужно вариировать гиперпараметры, изучать датасет, менять длительность обучения, топологию нейросети (больше слоёв, включать или выключать батч-нормализацию и т.д.), размер батчей или оптимизатор, снять барьер преждевременной остановки или же просто поменять настройки генератора случайных чисел.