Генерация изображений с моделями диффузии с использованием Keras и TensorFlow

Генерация изображений с моделями диффузии с использованием Keras и TensorFlow

Вы наверняка слышали о Dall-E 2. Опубликовано Open AI, которая представляет собой модель, которая генерирует реалистично выглядящие изображения из данной текстовой подсказки. Ознакомиться с уменьшенной версией модели можно здесь.

Вы когда-нибудь задумывались, как это работает под капотом? Ну… он использует новый класс генеративной техники, называемый «диффузией». Идея была предложена Золь-Дикштейном и соавторами в 2015 году. По сути, модель генерирует изображение из Шума.

Но зачем использовать диффузионные модели, когда вокруг есть GAN?

GAN отлично подходят для создания изображений высокой точности. Но, как указано в этой статье Open AI: Модели диффузии превосходят GAN в синтезе изображений, модели диффузии намного лучше синтезируют изображения, поскольку они более точно соответствуют изображению. GAN должны создавать изображение за один раз и, как правило, не имеют каких-либо опций для уточнения во время создания изображения. С другой стороны, диффузия — это медленный и повторяющийся процесс, во время которого шум шаг за шагом преобразуется в изображение. Это позволяет моделям диффузии иметь лучшие возможности для направления изображения к желаемому результату.

В этой статье мы рассмотрим, как создать собственную модель диффузии на основе вероятностных моделей диффузии с шумоподавлением (Ho et al, 2021) (DDPM) и неявных моделей диффузии с шумоподавлением (Song et al, 2021) (DDIM). с помощью Keras и TensorFlow. Итак, приступим…

Процесс, лежащий в основе моделей диффузии, делится на две части:
— процесс прямого шумоподавления и
— процесс обратного шумоподавления.

Шум вперед:

Концепция моделей диффузии основана на хорошо изученной концепции диффузии в физике.

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

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

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

Для процесса прямого зашумления q мы определяем цепь Маркова для предопределенного количества шагов, скажем, T. Который берет изображение и добавляет небольшое количество гауссовского шума к изображению в соответствии с графиком отклонения:β₀, β₁, … βt. Где β₀ ‹ β₁‹ … ‹ βt.

Затем мы обучаем модель, которая учится удалять это небольшое количество шума на каждом временном шаге (учитывая, что добавленный шум имеет небольшие приращения). Мы рассмотрим это в разделе обратного шумоподавления.

Но сначала, что такое цепь Маркова??

Цепь Маркова — это цепь событий, в которой событие определяется только предыдущим событием.

Здесь состояние x1 определяется только с помощью x0, x2 с помощью x1 и так далее, пока мы не достигнем xT. Таким образом, для наших целей состояние x0 является нашим нормальным изображением, и по мере продвижения вперед по нашей цепи Маркова изображение становится все более зашумленным, пока мы не достигнем состояния xT.

Добавление шума:

Согласно нашей цепи Маркова, состояние xt определяется только состоянием xt-1. Для этого нам нужно рассчитать вероятность q(xt|xt-1) создать слегка более зашумленное изображение на временном шаге t по сравнению с t-1. Это «слегка» зашумленное изображение создается путем выборки небольшого количества шума с использованием распределения Гаусса «N» и добавления его к изображению. Шум, выбранный из распределения Гаусса, определяется только средним значением и стандартным отклонением. Здесь мы используем график отклонений: β₀, β₁, … βt. Мы делаем среднее значение зависимым от βt и входного изображения xt. Итак, наконец, q(xt|xt-1)может быть определено как:

И в соответствии с принципом цепей Маркова вероятность возникновения цепи от x1 до xT для заданного начального состояния x0 определяется выражением:

Репараметризация:

Роль нашей модели состоит в том, чтобы отменить добавленный шум для каждой метки времени. Чтобы сгенерировать зашумленное изображение с указанной временной меткой, нам нужно повторять цепь Маркова, пока мы не получим желаемое зашумленное изображение. Этот процесс очень неэффективен. В качестве обходного пути мы используем прием перепараметризации, который использует аппроксимацию для создания шума в требуемую временную метку. Этот трюк работает, потому что сумма двух гауссовских выборок также является гауссовой выборкой. Вот формула репараметризации:

Следовательно, мы можем предварительно вычислить значения для α и α bar,используя формулу для q(xt|x0), получить зашумленную изображение xt на временном шаге t с исходным изображением x0.

Достаточно теории, давайте закодируем это…

Вот зависимости, которые нам понадобятся для построения нашей модели.

!pip install tensorflow
!pip install tensorflow_datasets
!pip install tensorflow_addons
!pip install einops

Начнем с импорта

Для этой реализации мы будем использовать набор данных цифр MNIST.

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

Теперь давайте визуализируем процесс прямого зашумления.

Обратное шумоподавление:

Давайте поймем, что именно будет делать наша модель.

Нам нужна модель генерации изображений, которая будет предсказывать, какой шум был добавлен к изображению в заданную временную метку. Эта модель должна принимать на вход зашумленное изображение вместе с отметкой времени и предсказывать, какой шум был добавлен к изображению на этом временном шаге. Модель в стиле U-Net идеально подходит для этой работы. Мы можем внести некоторые изменения в базовую архитектуру, заменив сверточные слои на уровни ResNet, добавить механизмы для учета кодирования временных меток, а также обратить внимание слои. Модель U-Net была впервые предложена для сегментации биомедицинских изображений, но с момента ее создания она была модифицирована и использовалась для множества различных приложений.

Позвольте кодировать нашу U-Net

1) Вспомогательные модули

2) Строительные блоки модели U-Net:
Здесь мы включаем встраивание времени путем масштабирования и сдвига входных данных, передаваемых в блок resnet. Этот коэффициент масштабирования и сдвига достигается путем передачи временных вложений через модуль Multi Layer Perceptron (MLP) в блоке resnet. Этот MLP преобразует временные вложения фиксированного размера в вектор, который соответствует совместимым размерам блоков на уровне resnet. Масштаб и сдвиг записываются как «Гамма» и «Бета» в приведенном ниже коде.

3) Модель U-Net

После того, как мы определили нашу модель U-Net, теперь мы можем создать ее экземпляр вместе с менеджером контрольных точек, чтобы сохранять контрольные точки во время обучения. Пока мы этим занимаемся, давайте также создадим наш оптимизатор. Мы будем использовать оптимизатор Adam со скоростью обучения 1e-4.

Обучение нашей модели:

Шаг обратного шумоподавления для нашей модели определяется как p, где p равно:

Здесь мы хотим, чтобы наша модель, т. е. наша модель U-Net, предсказывала шум во входном изображении xt на заданном временном шаге t, по существу предсказывая значение µ(xt, t) и Σ(xt, t), то есть среднее значение и дисперсия для xt на временном шаге t. Мы вычисляем потери для прогнозируемого шума между прогнозируемым шумом Є_θи исходным шумом Єпо следующей формуле:

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

Для процесса обучения мы будем использовать следующий алгоритм:
1) Сгенерировать случайное число для генерации временных меток и шума.
2) Создать список случайных временных меток в соответствии с размером пакета
3) Пропустите входное изображение через процесс прямого зашумления вместе с временными метками.
4) Получите прогнозы из модели U-Net, используя зашумленное изображение и временные метки.
5) Рассчитайте потери между прогнозируемый шум и реальный шум.
6) Обновите обучаемые переменные в модели U-Net.
7) Повторите для всех обучающих пакетов.

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

Здесь xt — случайная выборка, которую мы пропускаем через нашу модель U-Net и получаем Є_θ, затем вычисляем xt-1 в соответствии с формула:

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

Теперь давайте создадим наш алгоритм обратного шумоподавления, используя подход DDPM.

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

Вот пример GIF, созданного с использованием алгоритма логического вывода DDPM:

Есть одна проблема с алгоритмом вывода, предложенным в документе DDPM. Процесс очень медленный, так как нам нужно пройти через все 200 временных шагов, чтобы получить результат. Чтобы ускорить этот процесс, в документе DDIM был предложен улучшенный цикл логического вывода. Давайте обсудим это..

ДДИМ:

В документе DDIM авторы предложили немарковский метод для процесса обратного шумоподавления, тем самым сняв ограничение, согласно которому порядок цепочки должен зависеть от предыдущего изображения. В документе предлагается изменить цель DDPM, сделав функцию потерь более общей:

Из этой функции потерь мы можем сделать вывод, что значение потерь зависит только от q(xt|x0), а не от вероятности соединения q(x1:T|x0). Наряду с этим авторы также предложили изучить другой подход к выводу, который не является марковским. Приближается сложная математика:

Вышеупомянутые изменения также делают немарковским прямой процесс, где σ контролирует стохастичность прямого процесса. Когда σ→0, мы достигаем случая, когда xt−1 становится известным и фиксированным. Для генеративного процесса с фиксированным априором pθ(xT)=N(0,I):

Наконец, формула для вывода дается:

Здесь, если положить σ=0 ∀ t, то прямой процесс становится детерминированным.
Приведенные выше формулы взяты из [1].

Хватит математики, давайте закодируем это.

Теперь давайте воспользуемся процессом обратного шумоподавления, аналогичным DDPM. Обратите внимание, что мы используем только 10 шагов для этого цикла вывода вместо 200 шагов DDPM.

Вот пример GIF из вывода ddim:

Эту модель также можно обучить на другом наборе данных, и код, приведенный в этом посте, достаточно надежен, чтобы поддерживать изображения с более высоким разрешением и RGB. Например, я обучил модель на наборе данных celebA генерировать изображения 64x64 rgb, вот некоторые из результатов:

На этом мы можем закончить эту тему. Существует много связанной литературы, основанной на концепции моделей диффузии. Вот несколько интересных материалов:
1) GLIDE: на пути к созданию и редактированию фотореалистичных изображений с помощью текстовых моделей диффузии
2) Сверхвысокое разрешение изображения с помощью итеративного уточнения
3) Диффузионные модели превзошли GAN по синтезу изображений
4) Imagen
5) Dall-E 2

[1] Изучение моделей распространения с помощью JAX, Даршан Дешпанде. "связь".

Вы также можете прочитать продолжение этой истории, где я рассказываю о том, как генерировать изображения из меток классов. "связь".

Что подключить? Пожалуйста, напишите мне на [email protected]