Построение нейронной сети раскрашивания изображения

Построение нейронной сети раскрашивания изображения. Часть 4. Реализация

Добро пожаловать в четвертую и последнюю часть этой серии, где мы, наконец, реализуем нейронную сеть, способную применять цвет к черно-белым изображениям. В предыдущих статьях мы рассмотрели основы генеративных моделей и автоэнкодеров, искусственных нейронных сетей и сверточных нейронных сетей. Если это звучит для вас тарабарщиной, обязательно ознакомьтесь с соответствующими статьями, прежде чем изучать следующую часть (ссылки ниже).

Вся серия состоит из следующих 4 частей:

Часть 1: описывает основы генеративных моделей и автоэнкодеров.
Часть 2: демонстрирует основные концепции искусственных нейронных сетей.
Часть 3: представляет базовые знания о свёрточных нейронных сетях.
Часть 4 (текущая): описывает реализацию фактической модели.

Отказ от ответственности: это ни в коем случае не учебник. Он дает некоторые элементарные знания, но основная цель — продемонстрировать, как можно построить такую модель.

Вся модель была построена с помощью PyTorch, а предварительная обработка изображения была выполнена с помощью библиотеки Scikit-Image. Весь код можно найти в: https://github.com/gkamtzir/cnn-image-colorization

Данные и предварительная обработка

Прежде чем приступить к фактической реализации, нам нужен довольно большой набор данных, содержащий раскрашенные изображения. Имейте в виду, что наш подход не требует соответствующих черно-белых изображений, потому что, как мы упоминали в первой статье, мы будем использовать формат LAB, что означает, что мы можем разложить изображения обучающей выборки и получить черно-белая версия каждого. Я выбрал набор данных Image Colorization Dataset, содержащий 5000 раскрашенных изображений для обучения и 739 изображений для тестирования. Размеры каждого изображения 400x400x3. Их содержание варьируется от изображений еды, людей, животных, транспортных средств до изображений экстерьера и интерьера.

Единственной предварительной обработкой было преобразование RGB в формат LAB. Для этой цели я использовал библиотеку Scikit-Image вместе с классом Dataset, который PyTorch предоставляет для создания механизма чтения и загрузки изображений. Дополнительные сведения см. в файле Dataset.py.

Архитектура и конфигурация

Касательно архитектуры нейронной сети мы уже упоминали, что попробуем реализовать автоэнкодер, где энкодер будет состоять из сверточных слоев, а декодер будет содержать транспонированные сверточные слои. На вход будут изображения 400x400 с 1 каналом, значением L (яркость). На выходе мы получим изображение размером 400x400 с двумя каналами, значениями a и b. Окончательное раскрашенное изображение будет построено путем объединения предсказанных a и b с входными данными L.

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

Одна вещь, о которой мы не говорили, — это нормализация пакетов между слоями. Я обещаю написать статью о пакетной нормализации, а пока вы можете думать о пакетной нормализации как об умном трюке, который позволяет нам ускорить процесс обучения, не забывая при этом о весовых значениях.

Пакетная нормализация — это гораздо больше. Я объясню все в следующей статье.

В качестве функции активации я использовал ReLU, так как это один из самых безопасных вариантов. Размер пакета был установлен на 32, что соответствует обучению сети с пакетами, содержащими 32 объединенных изображения. Потери прогноза рассчитывались по среднеквадратической ошибке или MSE, тогда как в экспериментах сеть использовала оптимизатор Адама.

Всего я экспериментировал с 6 разными сетями. В последующих разделах я предоставлю настройки и результаты для всех из них. Каждый из них обучался на наборе для разработки из 200 изображений примерно для 200 эпох. Наборы для разработки используются для тестирования и сравнения различных архитектур перед передачей всего набора данных в модель и могут сэкономить много времени.

1-я Архитектура

Первая и самая базовая архитектура показана на следующем изображении. Мы провели 3 отдельных эксперимента с 3 разными скоростями обучения, но все они работали примерно одинаково. Кроме того, они не смогли изучить даже обучающие экземпляры, что является тревожным сигналом с точки зрения возможностей обучения.

2-я Архитектура

Вторая архитектура, по сути, является усовершенствованием предыдущей, где я добавил 1 сверточный слой и 1 транспонированный сверточный слой в кодере и декодере соответственно. В этом случае потери были снижены по сравнению с первым случаем, а эксперимент со скоростью обучения 0,001 показал лучшие результаты. Тем не менее, сеть не могла раскрашивать изображения, хотя и пыталась раскрасить некоторые места.

3-я архитектура

Третья архитектура была не просто усовершенствованием, а полной модификацией. Я настроил сеть для реализации так называемой U-Net [1]. Схема U-Net использует ранее рассчитанные выходные данные кодера в качестве входных данных в следующих частях декодера. Таким образом, мы гарантируем, что сеть не потеряет важную информацию. Точную структуру U-Net можно найти ниже:

При таком подходе сеть сходилась быстрее с меньшими ошибками как на обучающих, так и на тестовых выборках. Кроме того, это была первая сеть, которая смогла последовательно применять определенные цвета.

4-я архитектура

Эта архитектура была основана на предыдущей с добавлением одного дополнительного уровня как к кодировщику, так и к декодеру. Конечные результаты показали еще более низкие потери в тренировочном наборе, а выходные данные содержали больше раскрашенных областей. Обратите внимание, что потери при тестировании были увеличены. Это означает, что наша модель превосходит тренировочный набор, но на данный момент это не проблема, потому что первый шаг при построении модели — убедиться, что она способна обучаться. Это достигается за счет того, что модель до некоторой степени переоснащается, а затем, увеличивая тренировочный набор, проблема переобучения обычно исчезает.

5-я Архитектура

На этом этапе я решил изменить макет сети, введя слои расширения, которые также известны как «слои trous». Есть исследования, указывающие на значительное улучшение прогнозирования в таких случаях, как наш [2]. Опять же, потери были еще больше уменьшены, и модель смогла более точно раскрашивать изображения.

6-я Архитектура

Последняя архитектура — это расширенная версия пятого корпуса, где я добавил 2 дополнительных слоя. Результаты не сильно отличались от предыдущей версии, что делает эту архитектуру хорошей точкой остановки. Справедливости ради, в некоторых местах архитектура 5 была лучше, чем 6, но поскольку последняя демонстрировала меньшее переоснащение, я выбрал ее в качестве окончательной модели.

Окончательные результаты

Обучение набору для разработки под 6-ю архитектуру длилось около 40 минут на графических процессорах Google Colab Pro и около 2,5 часов на процессоре i5–4690K@3,9ГГц. Из-за ограничений по времени и доступности графического процессора я был ограничен использованием только центрального процессора для обучения. Вот почему окончательная архитектура была обучена на 2000 изображений, а не на всем наборе данных. Итак, я тренировал модель на 300 эпох со скоростью обучения 0,001 в течение 3 дней. Окончательные результаты были обнадеживающими, так как модель смогла раскрасить не только изображения, с которыми она столкнулась во время обучения, но и изображения, которые она раньше не видела!

Обучающий набор

Набор для тестирования

И на этом завершается вся серия о решении задачи раскрашивания изображений с помощью нейросетей. Я очень надеюсь, что вы многому научились в процессе и в то же время весело провели время. Это была моя первая серия статей, посвященная машинному обучению, поэтому будьте уверены, что в ближайшем будущем их будет еще много. А пока продолжайте учиться!

Построение нейронной сети раскрашивания изображения — Часть 4: Реализация