Как стать автором
Обновить

Как сделать бота, который превращает фото в комикс. Часть вторая. Обучение модели

Время на прочтение5 мин
Количество просмотров11K
⇨ Первая часть
⇨ Третья часть

И снова здравствуйте!

Как вы могли заметить, праздники несколько подкосили график выхода статей.
Думаю, многие за это время успели если не полностью обучить свою модель, то хотя бы поэкспериментировать с различными наборами данных.
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}
Для того, чтобы не путаться, какая модель на каком датасете дает какой результат,
начните вести небольшой дневник. Записывайте исходные данные и итоги обучения.

Например, командную строку и эпоху, до которой вы дошли. Этих данных хватит, чтобы сохранить параметры обучения и датасет, на котором мы тренировались.

Ведь чем отличается научный эксперимент от валяния дурака? Тем, что у нас все задокументировано!
На этом все на сегодня! В следующей статье мы с вами научимся экспортировать готовую модель и захостим ее в облаке.

Не бойтесь экспериментировать. Обязательно попробуйте несколько разных наборов данных, сравните результаты и поделитесь ими в комментариях!

До новых встреч!

⇨ Следующая часть
Теги:
Хабы:
Всего голосов 29: ↑29 и ↓0+29
Комментарии4

Публикации

Истории

Работа

Data Scientist
97 вакансий
Python разработчик
199 вакансий

Ближайшие события

19 августа – 20 октября
RuCode.Финал. Чемпионат по алгоритмическому программированию и ИИ
МоскваНижний НовгородЕкатеринбургСтавропольНовосибрискКалининградПермьВладивостокЧитаКраснорскТомскИжевскПетрозаводскКазаньКурскТюменьВолгоградУфаМурманскБишкекСочиУльяновскСаратовИркутскДолгопрудныйОнлайн
3 – 18 октября
Kokoc Hackathon 2024
Онлайн
10 – 11 октября
HR IT & Team Lead конференция «Битва за IT-таланты»
МоскваОнлайн
25 октября
Конференция по росту продуктов EGC’24
МоскваОнлайн
7 – 8 ноября
Конференция byteoilgas_conf 2024
МоскваОнлайн
7 – 8 ноября
Конференция «Матемаркетинг»
МоскваОнлайн