Антон Белов

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

Первая мысль - использовать ChatGPT и его встроенные инструменты. Но как бы он ни обещал отличный результат, то ресурсов не хватает, то он подозревает меня в неправомерном использовании фотографий. Можно было попробовать другой AI сервис или вовсе найти специализированное решение, но меня быстро догнала вторая мысль.

Вторая мысль - почему бы не собрать свое? Задача кажется простой: переложить стиль с одной картинки на другую. На выходе получим in-house решение, и чуть больше понимания, как оно работает “под капотом”. Короче говоря, я открыл IDE и начал писать на питоне, вспоминая университетскую молодость.

Что имеем на руках

Собственно, (кроме энтузиазма) имеется следующее:

▪Style image - изображение, с которого забираем стиль
▪Content image - изображение, на которое переносим стиль
▪Output image - так как есть референсы с используемого ранее веб-приложения, то есть ожидание, как должен выглядеть результат
▪Apple M1 Pro в качестве станка

Так как это типо ML (машин лернинг), то питон - обязателен. Но, как потом оказалось, каких-то прям навыков в нем эта задача не требует. Тем более, когда есть инструменты типа Cursor AI.

Основные подходы

В машинном обучении есть несколько подходов к решению задачи стилизации изображений. Вот основные из них:

▪Классический Neural Style Transfer (NST) использует глубокие сверточные нейросети (CNN, например, VGG-19) для раздельного выделения контента и стиля изображения. Затем с помощью оптимизации (L-BFGS или Adam) создается новое изображение, минимизирующее разницу со стилем.

▪Fast Style Transfer - вместо оптимизации обучается CNN, которая мгновенно применяет стиль к любому изображению. Для этого создается специальная генеративная сеть (Feed-forward CNN), предобученная на конкретном стиле.

▪Arbitrary Style Transfer использует Adaptive Instance Normalization (AdaIN) для адаптации стиля к изображению на лету, без обучения новой модели.

▪Diffusion Models - генеративные модели вроде Stable Diffusion позволяют создавать изображения в стиле референсов, сочетая их с текстовыми подсказками, масками, глубиной сцены и другими условиями. Тут будущее (DALL-E, например).

Много сложных слов и аббревиатур. Но список упрощается до одной сравнительной таблички:

Для начинающего ML инженера отлично подходит классический NST, так как не надо ничего часами обучать на GPU или обладать знаниями по генеративно-состязательным сетям. Кроме того этот подход отлично годится для экспериментов на старте.

Немного теории классического NST

В классическом NST, предложенном еще в 2015 (тут статья), используется предобученная VGG-19, которая позволяет выделить контент и стиль изображения. VGG-19 - это глубокая сверточная нейросеть (CNN), разработанная Oxford Visual Geometry Group (VGG). Она была обучена на огромном датасете изображений ImageNet для распознавания объектов.

Основные этапы реализации NST следующие:

1.Загружаем изображения (контентное и стилевое)
2.Пропускаем их через VGG-19 и извлекаем признаки контента и стиля
3.Генерируем изображение (из копии контента или случайного шума)
4.Вычисляем общую функцию потерь
5.Обновляем пиксели, используя градиентный спуск
6.Повторяем процесс 300-1000 раз, пока изображение не станет стилизованным

VGG-19 состоит из двух частей: сверточные извлекают признаки изображения, полносвязные делают классификацию. NST использует только сверточные слои, потому что нужны признаки контента и стиля, а не классификация. Причем только некоторые слои. Модель используется как экстрактор, ее веса не будут изменяться.

Функция потерь в NST отвечает за комбинирование контента одного изображения со стилем другого и состоит из двух частей:

▪Content Loss - Задача: сохранить структуру объектов
▪Style Loss - Задача: перенести текстуру, цвет и мазки кисти

Дополнительно еще используют Total Variation Loss для сглаживания изображения и устранения мелких артефактов.

Градиентный спуск - оптимизационный алгоритм, который минимизирует функцию потерь - итеративно обновляет пиксели изображения. Два наиболее популярных оптимизатора для NST - L-BFGS и Adam. Таким образом, NST - это оптимизация пикселей, а не весов модели!

Немного про реализацию

Для реализации потребуется Python (версии 3.9 и новее), torch - ядро фреймворка PyTorch для обучения и torchvision - дополнение для работы с изображениями.

Рассказывать про написание не буду, так как примеров и материалов много. Исходный код нужно найти тут, приведу используемые параметры:

С L-BFGS на 150 итерациях и Adam на 1500 получилось добиться приемлемого результата.

Для своей реализации можно форкнуть готовое (оно на Lua), мне же просто хотелось разобраться, пройдя по граблям. Дальше будет конспект про опыты на Боромирах.

Альфа и бета

В NST параметры α (альфа) и β (бета) управляют балансом между контентом и стилем в итоговом изображении. Иными словами, α - это вес Content loss, а β - вес Style loss.

Абсолютные значения не важны, имеет значение отношение α / β, поэтому можно выставить α = 1 и менять только β. Эффект переноса деталей из стиля начинает замечаться при достаточно большом β, начиная от 1e4. При слишком высоком значении β изображение постепенно теряет форму оригинального содержимого (но делает это красиво).

Значение 1e6 оказалось вполне балансным.

Вклад стилевых признаков

При вычислении потери стиля можно управлять вкладом стилевых признаков из разных слоев сверточной нейросети. Каждый слой сети отвечает за разное:

▪Ранние слои (conv1_1, conv2_1) - про мелкие текстуры, цвета, мазки кисти
▪Средние слои (conv3_1, conv4_1) - про более сложные узоры и формы
▪Глубокие слои (conv5_1) - про композиционные особенности стиля

Конечно, в качестве эксперимента тут же захотелось выставить β в 1e9 и обнулить веса на всех слоях кроме одного. Таким образом явно наблюдать эффект того или иного стилевого признака.

Можно сделать следующие выводы:

▪conv1_1 - Яркие текстурные эффекты; много резких мазков и цветовых артефактов; контуры исказились. Слой сильно учитывает локальные текстуры и цвета, но не сохраняет форму объекта.

▪conv2_1 - Более сглаженный эффект; все еще присутствуют текстуры, но они более мягкие; контуры лучше сохранены, но стиль все равно заметно влияет. Слой балансирует текстуры и сохранение формы.

▪conv3_1 - Видны крупные мазки и узоры; контуры становятся более четкими; изображение похоже на нарисованное, но без хаотичных текстур. Слой учитывает не только текстуру, но и структурные элементы стиля, делая перенос более сбалансированным.

▪conv4_1 - Текстурные эффекты минимальны, хорошо сохранены общие контуры, стиль выражается через цветовые вариации и освещение. Слой больше отвечает за крупные формы и композицию, создавая мягкий стилизованный эффект.

▪conv5_1 - Очень слабый перенос стиля, есть легкий шум. Слой учитывает глобальные особенности, но не дает выразительных текстур.

Для переноса стиля с живописных картин можно использовать следующие настройки:

▪conv1_1 = 1.0 - Максимальный вес, потому что живопись имеет ярко выраженные текстуры и мазки.
▪conv2_1 = 0.75 - Чуть меньше, чтобы не перенести слишком резкие детали.
▪conv3_1, conv4_1, conv5_1 = 0.2 - Низкие веса, чтобы не разрушить форму изображения.

Разные оптимизаторы

Как уже было сказано, в классическом NST ключевая задача - это оптимизация пикселей изображения, которое мы генерируем. Тут два основных оптимизатора:

▪L-BFGS - быстро и точно сходится, потребляя меньше памяти, особенно эффективен при оптимизации изображения напрямую.

▪Adam - адаптивный оптимизатор с хорошей стабильностью, но требует больше памяти и итераций, удобен на GPU и гибок в использовании.

Гибкость Adam’а почти не проверял, задав рекомендуемые чатом GPT параметры с третьего раза. А вот про скорость сходимости - интересно. Тогда как L-BFGS выдает нормальный результат уже на 100-150 итераций, то для Adam надо 1500 и больше. Однако, скорость итерирования раз в двадцать выше, это на Apple M1 Pro.

Отличие по памяти легко измерить через psutil и получилось, что Adam в этой задаче тратит в четыре раза больше RAM. В теории L-BFGS должен тратить больше на больших моделях, но в этой задаче у него всего один параметр - одно изображение.

Говорят, что Adam может давать менее четкие результаты, но как по мне, так для стилизации его небольшие эффекты размытия мягкими мазками даже лучше.

Копия, шум и блендер

С входным изображением для оптимизатора можно проводить любопытные эксперименты. Например, это может быть копия оригинала, случайный шум (тензор той же размерности), или даже смесь - например, 90% от оригинала и 10% от стиля, объединенные в один тензор на входе.

При использовании шума имеет смысл увеличить α (альфа), поскольку контент теперь нужно восстанавливать с нуля. Конечно, это потребует и большего количества итераций.

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

Разные девайсы

Запускать все это можно на разных типах устройств, и в коде необходимо явно указывать, на каком именно типе устройства будет работать PyTorch. Разберем варианты и померим наконец перформанс:

▪CPU - есть на любом компьютере. Работает медленно, поскольку NST требует множество сверток и градиентных шагов, которые на CPU выполняются на порядки медленнее, особенно при больших изображениях.

▪MPS (Metal Performance Shaders, macOS GPU backend) - доступен на Apple Silicon (M1/M2/M3). Дает значительное ускорение по сравнению с CPU, но пока не такой гибкий и отлаженный, как CUDA.

▪CUDA (NVIDIA GPU backend) - используется на видеокартах NVIDIA. Самый быстрый и стабильный вариант для NST, особенно эффективен при использовании оптимизатора Adam, но это не точно.

▪TPU (Tensor Processing Unit) - предоставляется Google. Очень быстрый, но ориентирован на инференс и обучение моделей, а не на оптимизацию входного тензора; поддержка в PyTorch ограничена.

Доступ к CUDA и TPU можно получить через Google Colab - онлайн-сервис от Google, который позволяет запускать Python-код в браузере. На бесплатной подписке не так много compute units (доступных ресурсов) и не такое мощное железо, но поиграться точно хватит.

Коротко о результатах запуска NST на разных устройствах:

▪На CPU вообще нет смысла запускать, можно состариться.
▪Интегрированный GPU на Apple M1 Pro показал себя хорошо, и в полтора-два раза хуже чем Colab T4 GPU - это примерно как сравнивать с RTX 2060.
▪Как и ожидалось, на TPU (Colab v2-8) были проблемы: L-BFGS не завелся вообще. Но Adam пробежал даже быстрее чем на бесплатном CUDA.

Резюме, что дальше?

Получилось подразобраться для себя и для бумаги, хотя бы на уровне пользователя. Классическому алгоритму NST уже около девяти лет, а вот попробовать по-DeepResearh’ить его параметры и коэффициенты, позапускать на обилии разных устройств, включая облака - почему бы и не сейчас, и очень интересно.

И это только первый шаг! Потому что ждать по 10 минут генерацию одной картинки - это не дело. Так что дальше - Fast Style Transfer; лишь бы хватило compute units уже не на оптимизацию пикселей, а на обучение нейросети-генератора. Но это уже совсем другая история.

Исходный код тут