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

  • Tutorial
⇨ Первая часть
⇨ Третья часть

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

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

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

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

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

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

⇨ Следующая часть

Средняя зарплата в IT

120 000 ₽/мес.
Средняя зарплата по всем IT-специализациям на основании 3 897 анкет, за 1-ое пол. 2021 года Узнать свою зарплату
Реклама
AdBlock похитил этот баннер, но баннеры не зубы — отрастут

Подробнее

Комментарии 4

    +1
    Спасибо за статью!

    1) Вы как-нибудь отслеживали прогресс? (tensorboard, консолька)

    2) Идея: получив комиксы от CycleGAN, вручную в фотошопе убрать артефакты, и скормить получившиеся пары оригинал/комикс уже в pix2pix.
      +1
      1) Отслеживание прогресса в этом дистрибутиве сделано через visdom, очень наглядным образом. В консоль выводятся показатели функций потери, но для GAN они не очень репрезентативны (только если что-то пошло ну уж совсем не так)
      2) Спасибо за идею! Еще можно попробовать существующие экшены\фильтры для фотошопа
        +1
        1) Ограничение на размер картинки можно обойти с помощью полносвёрточных нейросетей. Их можно обучать на 256х256, а потом применять хоть к 8K (на сколько оперативки хватит). Вот тут, например, так сделана трансформирующая нейросеть (но это не GAN).
        2) Свежий stylegan2 от NVIDIA даёт впечатляющие результаты.
        Сейчас делаю совмещение этих двух подходов. А именно — сначала научить дискриминатор с помощью GAN-а, а потом этим дискриминатором учить трансформирующую нейросеть. Сохранение содержания пока планирую считать через VGG.
          0
          Спасибо за интересный комментарий!
          1) Тут ограничение на размер картинки упирается в оперативку как при обучении, так и при использовании. Просто при обучении у нас 4 модели, а на проде — 1. И генераторы полносверточные, кстати (ну, они везде такие обычно)

          2) Есть подход NoGAN+GAN, когда обучают генератор отдельно, используя VGG в качестве функции потери, генерят фейковые изображения, дальше обучают критик на фейковых и подлинных изображениях (как обычный классификатор, без GAN). Затем берут эти модели обучают как обычную GAN, только не с нуля, что позволяет сократить время обучения и избежать заведомо плохих результатов на начальном этапе обучения.
          У меня есть модель, обученная по принципу NoGAN + GAN. Бот на ее основе: t.me/@pimpmyresbot
          (увеличивает разрешение картинки в два раза, до максимум в 1500х1500 пикселей)

        Только полноправные пользователи могут оставлять комментарии. Войдите, пожалуйста.

        Самое читаемое