Как сделать бота, который превращает фото в комикс. Часть вторая. Обучение модели
⇨ Первая часть
И снова здравствуйте!
Как вы могли заметить, праздники несколько подкосили график выхода статей.
Думаю, многие за это время успели если не полностью обучить свою модель, то хотя бы поэкспериментировать с различными наборами данных.
1. Ставим дистрибутив
2. Качаем фотки
3. ???
4. Profit!
Если же вам было не до этих наших нейросетей, или вы начинаете чтение с этой статьи, то, как говорится, нет времени объяснять, берем дистрибутив, качаем нужные фотки, и поехали!
Подробная инструкция — в предыдущей статье цикла.
Первый блин комом
Когда я обучал модель для @photo2comicsbot в первый раз, я, не мудрствуя лукаво, просто запихнул около 1000 страниц комиксов в датасет.
Да, вместе с обложками, анонсами и прочими филлерами.
На входе выглядело это вот так:
Результат получился соответствующий:
Модель отлично передает общие отличия между комиксами и фото:
- Характерная цветокоррекция
- Выделение контуров
- Разбивка на блоки
- Облачка с текстом
В принципе, этого варианта может быть вполне достаточно для продакшена.
Но мне было интересно, какие результаты можно получить на более «чистом» датасете.
Дубль два
Для следующего эксперимента я оставил страницы с одним блоком, а из остальных повырезал наиболее крупные куски, состоящие из одного блока, дабы избавиться от вертикальных и горизонтальных разделительных линий в итоговой модели.
Если вы вносите изменения в свой датасет, например, удаляете или добавляете фотографии, не забывайте сохранить предыдущий вариант. Вам будет куда вернуться, если что-то пойдет не такю
Второй датасет выглядел так:
Результат, как говорится, налицо:
Разделительные линии пропали, но вместе с ними ушла выразительность в облачках с текстом: они стали встречаться гораздо реже в датасете и перестали быть основной фичей.
Пустое пространство, которые раньше заполнялось текстом, теперь заполняется психоделическим шумом.
Third time’s a charm
Я решил избавиться от текста полностью, оставив только те части комиксов, где его нет. Не выжил никто: ни облачка, ни заголовки, ни названия серий.
Датасет после «лингвистического геноцида»:
Итоговый результат:
С одной стороны, артефактов стало меньше, с другой — ушла выразительность. Поскольку в обучении GAN качество результата в большей степени определяем мы с вами, нет единого рецепта хорошей модели.
Красота в глазах смотрящего, поэтому не бойтесь экспериментировать и выбирайте то, что больше нравится.
К оружию!
Ну что ж, об основных подводных камнях и методах борьбы с ними мы уже поговорили, пора переходить к делу.
Переходим в папку, куда мы скачали дистрибутив. Напомню, она называлась pytorch-CycleGAN-and-pix2pix
Поскольку инструкция у нас для чайников, мы не будем трогать код дистрибутива, благо все необходимые параметры можно задать из командной строки.
В данном дистрибутиве обучение по умолчанию идет в течение двухсот эпох, с линейным затуханием скорости обучения (learning rate) после сотой эпохи.
Максимальное разрешение, которое смогла осилить моя восьмигигабайтная карточка, — 400×400. Полное обучение заняло у меня около 33 часов, поэтому мы с вами воспользуемся небольшим лайфхаком. Сначала мы обучим модель на картинках 128×128, затем — 256×256, и только на финальном этапе покажем ей наши шикарные 400×400.
Первый этап продлится 100 эпох, второй и третий — по 50, что позволит нам сократить время обучения практически вдвое.Как показывает практика, результат при таком подходе не хуже, а иногда и лучше, чем при обучении сразу на максимальном разрешении.
Пожалуй, хватит теории, пора переходить к практике.
В командной строке вводим.
Этап первый
python train.py --dataroot {dataset root folder} --name {model name} --model cycle_gan --crop_size 128 --load_size 180 --init_type kaiming --netG resnet_9blocks --no_dropout --batch_size 4
Не забываем заменить фигурные скобки на свои значения.
Разберем некоторые параметры подробнее:
--batch_size {number}
Этот параметр отвечает за количество обрабатываемых картинок «за такт», положительно влияя на скорость и отрицательно — на прожорливость модели.
На каждом этапе подбирайте максимально возможный --batch_size, который не вызывает ошибок из-за нехватки памяти GPU.
-- dataroot {dataset root folder}
— папка с нашим датасетом. Внутри нее должны находиться папки trainA, trainB, testA, testB, как описано в предыдущей статье.
--name {model name}
— название вашего проекта. Может быть произвольным, но я рекомендую включить в название архитектуру модели, максимальное разрешение и имя датасета.
Например: «resnet9_128to400_comics8»
Так вы сможете различить эксперименты с различными параметрами и данными.
Первый этап обучения можно прекратить на сотой эпохе.
За прогрессом можно наблюдать из браузера: localhost:8097/
(либо другая ссылка, которая будет видна в консоли)
Этап второй
python train.py --dataroot {dataset root folder} --name {model name} --model cycle_gan --crop_size 256 --load_size 290 --init_type kaiming --netG resnet_9blocks --no_dropout --batch_size 2 --epoch 100 --epoch_count 0 --continue_train
Здесь важно указать ту эпоху, на которой мы закончили обучение в первом этапе.
»--epoch 100» означает что мы загрузим модель из чекпоинта сотой эпохи
»--epoch_count 0» означает, что мы начнем обучение с нуля, с максимальной скоростью обучения.
Второй этап обучения можно прекратить на 50-й эпохе.
Этап третий
python train.py --dataroot {dataset root folder} --name {model name} --model cycle_gan --crop_size 400 --load_size 430 --init_type kaiming --netG resnet_9blocks --no_dropout --batch_size 1 --epoch 50 --epoch_count 0 --continue_trainТретий этап обучения можно прекратить на 50-й эпохе, но это дело вкуса. Можете доучить до конца, и выбрать понравившийся промежуточный результат. Важно помнить, что результат на 200-й эпохе может быть хуже, чем на 150-й.
А теперь выйди и посмотри, что ты наделал
В процессе обучения модель и промежуточные результаты будут сохранены в папку
/pytorch-CycleGAN-and-pix2pix/checkpoints/{model name}
Чтобы проверить модель, достаточно ввести в командную строку:
python test.py --dataroot {dataset root folder} --name {model name} --model cycle_gan --netG resnet_9blocks --crop_size 512 --load_size 580 --epoch {epoch name}
Вы можете посмотреть результат на тестовом датасете для любого чекпоинта, просто укажите его в качестве {epoch name}. Если {epoch name} не указан, будет взят последний чекпоинт.
Результат будет сохранен в папке: /pytorch-CycleGAN-and-pix2pix/results/{model name}/test_{epoch name}
Для того, чтобы не путаться, какая модель на каком датасете дает какой результат,
начните вести небольшой дневник. Записывайте исходные данные и итоги обучения.Например, командную строку и эпоху, до которой вы дошли. Этих данных хватит, чтобы сохранить параметры обучения и датасет, на котором мы тренировались.
Ведь чем отличается научный эксперимент от валяния дурака? Тем, что нас все задокументировано!
На этом все на сегодня! В следующей статье мы с вами научимся экспортировать готовую модель и захостим ее в облаке.
Не бойтесь экспериментировать. Обязательно попробуйте несколько разных наборов данных, сравните результаты и поделитесь ими в комментариях!
До новых встреч!
