С момента своего появления в 2013 через эту статью вариационный автокодировщик (VAE) как тип генеративной модели штурмовал мир байесовского глубокого обучения с его применением в широком диапазоне областей. Оригинальная статья Кингмы и Веллинга имеет более 10 тысяч цитирований; Между тем, поскольку на первый взгляд ее конструкция может показаться непростой для восприятия, было опубликовано множество замечательных статей, объясняющих интуицию, архитектуру и другие различные компоненты модели.

Однако реализация VAE обычно является дополнением к этим статьям, а о самом коде говорят меньше, особенно о том, что он контекстуализирован в какой-то конкретной библиотеке глубокого обучения (TensorFlow, PyTorch и т. Д.), Что означает, что код просто выкладывается в блоке кода, без достаточного количества комментариев о том, как работают одни аргументы, почему выбирают именно эту функцию по сравнению с другими и т. д. Кроме того, из-за гибкости каждой из этих популярных библиотек вы можете обнаружить, что все эти демонстрационные реализации VAE кажутся разными для каждой Другие. Более того, некоторые версии могут быть реализованы некорректно, даже с помощью одного из собственных руководств TensorFlow (я обращусь к этому сообщению позже, в основном в разделе Версия a. моей реализации), но ошибки могут не быть обнаружены. без сравнения с другими версиями. Наконец, конечно, приятно пройти через одну работающую версию, но я почти не встречал сообщения, в котором сравнивались бы разные способы реализации. Все это мотивирует меня написать этот пост.

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

Я планирую разделить этот пост на два компонента, которые вместе завершают заголовок:

  1. На что следует обратить особое внимание при внедрении VAE в целом с использованием TF 2 и TFP
  2. Как реализовать VAE по-разному, используя TF 2 и TFP

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

Я думал, что хорошо разбираюсь в модели после прочтения статьи, пока не начал ее реализовывать. При чтении статьи обычно бывает так много деталей о модели, что очень вероятно потерять из виду детали, на которые нужно обратить особое внимание. Иногда компонент модели, на которую в статье было потрачено одно короткое предложение, может занять у меня часы, чтобы полностью понять и заставить ее работать во время реализации. На мой взгляд, если вы хотите иметь хорошее представление о том, как работает модель, хотя определенно приятно попытаться изучить документ как можно более тщательно, ИМХО лучше попытаться реализовать это самостоятельно.

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

Часть 0: Уточнение цели реализации

В соответствии с обсуждением из приведенного выше раздела я представлю набор данных, а также конкретные задачи, которые должна выполнить модель VAE. Для этого поста я применил набор данных MNIST с рукописными цифрами с изображениями формы (28,28,1). Я предварительно обработал его, нормализовав набор данных, чтобы он находился между 0 и 1, и дискретизировал значения, чтобы они были либо 0, либо 1 , используя 0,5 в качестве порогового значения. Задачи модели VAE для этого набора данных:

(а) Реконструкция изображений входных цифр как можно ближе

(b) Генерация новых цифровых изображений, которые выглядят реалистично, с использованием случайных выборок из предыдущего распределения (а не выборок из апостериорного, обусловленного данными) в качестве входных данных для декодера.

Задачу (a) относительно легко выполнить, однако задача (b), на мой взгляд, та, которая указывает, работает ли реализация: что она усвоила мультимодальность различных цифр из обучающих данных и может создавать новые цифровые изображения в качестве генеративной модели.

Теперь мы готовы погрузиться в каждую составляющую этого поста.

Часть I: Основное внимание при внедрении VAE

Две вещи, обе относительно функция потерь для каждого экземпляра:

  • Функция потерь состоит из двух частей: обратного расхождения KL (между априорным и апостериорным распределением скрытой переменной z) и ожидаемого отрицательного логарифмического правдоподобия (распределения декодера для данных, подлежащих восстановлению. ; также называется ожидаемой ошибкой восстановления). Сложив эти две части вместе, мы получим отрицательный ELBO, который необходимо минимизировать. Для вычисления каждой части нам нужно обратить особое внимание на то, какую операцию (взятие суммы или взятие среднего) нам нужно выполнить с каждым измерением данных.
  • Вес KL-дивергенции в функции потерь - это гиперпараметр, который мы не должны игнорировать вообще: он регулирует «расстояние» между априорным и апостериорным распределением z и играет решающую роль в производительности модель.

Вот функция потерь:

Вычисление этой функции потерь может быть выполнено различными способами, и часто очень часто допускаются ошибки во время реализации, особенно в отношении w.r.t. что упомянуто в первом пункте: операции с каждым измерением данных. Я подробно остановлюсь на этом в части II, где каждый способ реализации представлен отдельно.

Обратите внимание, что это должна быть функция потерь для Beta-VAE, в которой ω может принимать значения, отличные от 1. Этот гиперпараметр имеет решающее значение, особенно когда для задачи (b), упомянутой в части 0: этот гиперпараметр определяет, насколько сильно мы хотим компенсировать разницу между предыдущими и апостериорное распределение z. Во многих случаях реконструкция изображений выглядела бы идеально, в то время как генерация из кода z, взятого из предыдущей версии, имела бы все виды сумасшедшего вида. Если ω слишком мало, мы в основном не регуляризуем апостериорирующую часть, поэтому после обучения она может сильно отличаться от предыдущей - в том смысле, что z отобранные из апостериорного распределения часто попадали в область с очень низкой плотностью апостериорного распределения. В результате декодер не будет знать, что делать с такими выборками z, поскольку он обучен на z из апостериорного распределения (обратите внимание на распределение в нижнем индексе математическое ожидание отрицательного члена логарифмической вероятности из функции потерь выше). Вот восстановленные и сгенерированные цифры при ω = 0,0001:

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

С другой стороны, если ω слишком велико, апостериорная часть будет подтянута слишком близко к предыдущей, поэтому независимо от того, какое изображение вводится в кодировщик, мы в конечном итоге получим z, как если бы он был выбран случайным образом из предыдущего. Вот восстановленные и сгенерированные цифры при ω = 20:

Обратите внимание, что все цифры выглядят одинаково, независимо от того, реконструирован ли он условно для входной цифры или сгенерирован как новая цифра. Это пример чрезмерной регуляризации.

В-третьих, у нас есть сценарий, когда ω установлен на «только правильную величину» - в том смысле, что апостериор достаточно отличается от предшествующего, чтобы можно было гибко обусловливать входные цифровые данные, таким образом, реконструкция выглядит великолепно; в то время как области с высокой плотностью двух распределений имеют достаточное перекрытие, поэтому образцы z из предыдущего не будут выглядеть слишком незнакомыми для декодера в качестве его входных данных. Вот восстановленные и сгенерированные цифры при ω = 3:

Обратите внимание, что все реконструированные и большинство сгенерированных цифр кажутся узнаваемыми, в то время как некоторые из сгенерированных цифр кажутся менее реалистичными.

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

Небольшое отступление: TensorFlow 2 и TensorFlow Probability

Вы могли заметить, что хотя я поместил эти две библиотеки в свой заголовок, пост до сих пор был сосредоточен исключительно на обсуждении модели VAE и ее применения в наборе данных MNIST. Изначально я хотел написать о TF 2 и TFP в качестве третьего компонента, когда я структурировал пост, но затем я решил контекстуализировать их, то есть поговорить о них пока я иду с помощью деталей реализации в части I и части II. Однако это может показаться слишком поспешным, если я сразу перейду к коду, не рассказывая немного об этих библиотеках. Итак, вот они.

TF2

TensorFlow 2.0.0 был выпущен в конце сентября 2019–, так что с момента его первого выпуска не прошло и года. Последняя версия на данный момент - 2.3.0, выпущенная в июле 2020. Я начал работать над исследовательским проектом, который заставил меня выбрать TF 2 в ноябре прошлого года, поэтому до сих пор я использовал все четыре версии, каждую некоторое время. В то время я был полностью осведомлен о том, что PyTorch набирает популярность в исследовательском сообществе; Я выбрал TF 2 в основном потому, что это библиотека, на которой была построена кодовая база проекта.

Мое первое впечатление от TF 2: действительно было намного удобнее разрабатывать модели глубокого обучения, чем TF 1. Поскольку нетерпеливое выполнение является одной из его наиболее отличительных особенностей, отделяющих TF 2 от TF 1, мне больше не нужно строить весь вычислительный граф, не видя каких-либо промежуточных результатов, что делает устранение неисправности, поскольку нет простого способа разложить каждый шаг и проверить его результат. Хотя нетерпеливое выполнение может привести к снижению скорости обучения модели, декоратор @tf.function помогает до некоторой степени восстановить эффективность в графическом режиме. Кроме того, его интеграция с Keras принесла несколько замечательных преимуществ - Последовательный и Функциональный API обеспечили разные уровни гибкости при наложении слоев нейронной сети (NN), точно так же, как их эквиваленты в PyTorch, torch.nn.Sequential и torch.nn.functional; также как высокоуровневый API, его интерфейс выглядит обманчиво простым (до такой степени, что это может вызвать проблемы во время реализации VAE - см. мое обсуждение в разделе реализации Версия a. в части II ), как будто я обучаю модель scikit-learn.

Между тем, поскольку с момента его первого выпуска прошло меньше года, он определенно не так стабилен, как я бы надеялся, от популярной библиотеки. Я все еще помнил, что застрял на какой-то простой процедуре тензорной индексации: я реализовал одну и ту же функциональность в нескольких версиях, что привело меня к одному и тому же сообщению об ошибке; однако это был настолько простой шаг, что никто не ожидал, что это именно то место, которое вызывает ошибку. Оказалось, что после обновления версии TF до 2.1.0 после ее выпуска в начале января этого года модель работала без изменения кода. Самый последний пример: при добавлении слоя Dense после выравнивания тензора я получил сообщение об ошибке, касающееся измерения, которое также исчезло при обновлении версии TF с 2.2.0 до 2.3.0 .

Более того, его сообщество пользователей и разработчиков оказалось не так активно, как я ожидал. Я разместил несколько вопросов на странице проблем с Keras Github, в том числе упомянутый выше - все закончилось тем, что я ответил на свой вопрос и закрыл проблему. Некоторые проблемы получили комментарии от пользователей через несколько месяцев, но команда TensorFlow / Keras не обратила на них внимания.

Наконец, некоторая его документация не была простой и недостаточно организованной, чтобы следовать ей. Я потратил много времени, пытаясь понять, как на самом деле работает декоратор @tf.function, и в конце концов пришел к выводу, что лучший способ - просто попытаться имитировать рабочие примеры, не слишком заботясь в данный момент о логическом обосновании. Некоторые из его руководств также оказались небрежными или даже вводящими в заблуждение (из-за неправильной реализации) - примеры я приведу позже.

TFP

TensorFlow Probability была представлена ​​в первой половине 2018 как библиотека, разработанная специально для вероятностного моделирования. Он реализует трюк репараметризации под капотом, который обеспечивает обратное распространение для обучения вероятностных моделей. Вы можете найти хорошую демонстрацию уловки повторной параметризации как в статье VAE, так и в этой статье, в которой предложен алгоритм Байеса с помощью Backprop - в первой работе есть скрытые узлы для скрытой переменной z и выходные узлы декодера являются вероятностными, в то время как последний имеет обучаемые параметры (веса и смещения каждого слоя NN) как вероятностные.

С TFP нам больше не нужно явно определять среднее значение и параметр дисперсии для апостериорного распределения z, а также вычислять дивергенцию KL, что значительно упрощает код. Фактически, это может сделать реализацию слишком простой, чтобы можно было сделать это, не имея хорошего понимания VAE, потому что его основные компоненты, которые потребуют хорошего понимания модели, в основном все абстрагируются с помощью TFP. В результате можно легко допустить ошибки при использовании TFP для реализации VAE.

Я бы посоветовал вам начать с явной реализации трюка с повторной параметризацией и определения термина KL, если вы впервые внедряете VAE и хотите хорошо знать, как работает модель - я начну с этого способа реализации в качестве первого. версия в части II.

Часть II: Различные способы реализации VAE

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

Различные способы реализации VAE обусловлены различными вариантами, которые у нас есть для каждого из следующих 3 модулей:

  • выходной слой кодировщика
  • выходной слой декодера
  • функция потерь (в основном та часть, которая вычисляет ожидаемую ошибку восстановления)

Каждый модуль имеет следующие параметры:

выходной слой кодировщика

  1. tfpl.IndependentNormal
  2. tfkl.Dense, который выводит (объединенное) среднее и (необработанное) стандартное отклонение апостериорного распределения z

выходной слой декодера

  1. tfpl.IndependentBernoulli
  2. tfpl.IndependentBernoulli.mean
  3. tfkl.Conv2DTranspose, который выводит логиты

функция потерь

  1. negative_log_likelihood = lambda x, rv_x: -rv_x.log_prob(x)
  2. tf.nn.sigmoid_cross_entropy_with_logits
  3. tfk.losses.BinaryCrossentropy
  4. tf.nn.sigmoid_cross_entropy_with_logits + tfkl.Layer.add_loss
  5. Явное вычисление в стиле PyTorch по эпохам до with tf.GradientTape() as tape + tf.nn.sigmoid_cross_entropy_with_logits + tfkl.Layer.add_loss
  6. Явное вычисление в стиле PyTorch по эпохам до with tf.GradientTape() as tape + tf.nn.sigmoid_cross_entropy_with_logits

Боковое примечание:

у нас есть следующие сокращения модулей:

импортировать тензорный поток как tf
импортировать тензорный поток как tfp
tfd = tfp.distributions
tfpl = tfp.layers
tfk = tf.keras
tfkl = tf.keras.layers

Также обратите внимание, что архитектура NN кодера и декодера менее важна - я использую ту же архитектуру, что и это руководство по TF для всех версий моей реализации.

Можно подумать, что, поскольку вычисление расходимости KL может быть выполнено иначе (аналитическое решение против приближения MC), оно также должно быть модулем; но поскольку способ реализации этого термина более или менее определяется выбором выходного уровня декодера, и добавление его, поскольку четвертый модуль может чрезмерно усложнить представление, я предпочитаю говорить об этом в рамках одной из версий реализации ( Версия а.).

Реализации - это просто комбинаторика различных опций каждого из 3 модулей:

Теперь давайте углубимся в каждую версию реализации.

Версия a.

Эта версия, пожалуй, наиболее широко принятая - насколько мне известно, все мои коллеги-исследователи писали таким образом: кодировщик выводит узлы, которые представляют среднее значение и (некоторое преобразование, с диапазоном в) стандартное отклонение апостериорного распределения скрытой переменной z. Затем мы используем трюк с повторной параметризацией для выборки z следующим образом (из уравнения (10) из статьи VAE):

в котором ϵ выбирается из стандартного многомерного распределения Гаусса, а среднее значение и стандартное отклонение апостериорного распределения q выводятся детерминированно из кодировщик. Полученный z затем используется в качестве входных данных для декодера. Поскольку после предварительной обработки наши данные изображения x являются двоичными, естественно предположить многомерное распределение Бернулли (где все пиксели независимы друг от друга), параметры которого являются выходными данными декодера. - это сценарий, описанный в Приложении C.1 Bernoulli MLP как декодер в документе VAE. Следовательно, логарифмическая вероятность распределения декодера имеет следующую форму (Уравнение (11) статьи VAE):

где y - это выходные данные декодера в виде параметров Бернулли для каждого пикселя, а D - количество пикселей для каждого экземпляра. z здесь представляет собой одиночный отсчет от кодировщика; поскольку ожидаемый член отрицательного логарифма правдоподобия в функции потерь не может быть вычислен аналитически, мы используем приближение MC следующим образом (из уравнения (10) статьи VAE):

в котором термин z (с точно такой же формой надстрочного индекса, что и в приведенном выше уравнении с перепараметризацией) представляет l -й MC образец для экземпляра изображения i -й цифры. На практике обычно достаточно установить L = 1, т.е. требуется только одна выборка MC.

Другой член в функции потерь, обратная дивергенция KL, также может быть аппроксимирован с помощью выборок MC. Поскольку мы предполагаем, что распределение кодировщика q является многомерным гауссовским, мы можем напрямую подставить среднее значение и дисперсию, выведенные кодировщиком, в его функцию плотности:

в котором f представляет многомерную гауссову плотность. Точно так же мы можем установить L = 1 на практике.

Кроме того, если мы позволим q быть многомерным гауссовым с диагональной ковариационной матрицей, член дивергенции KL может быть вычислен аналитически, как показано в Приложении B статьи VAE:

где J представляет количество измерений z.

В этой версии реализации я помещаю код, который вычисляет стоимость мини-пакетных данных изображения (обратите внимание, что функция loss вычисляется для одного экземпляра; cost (функция представляет собой среднее значение потерь для всех экземпляров) в функцию vae_cost и определите шаг оптимизации для каждой эпохи с помощью функции train_step. Вот как они реализованы в TF 2:

Здесь необходимо учесть несколько вещей:

  • Вычисление ожидаемой ошибки реконструкции

Поскольку отрицательная логарифмическая вероятность распределения Бернулли, по сути, является причиной потери кросс-энтропии (если вы не видите ее сразу, этот пост дает хороший обзор), мы можем использовать существующие функции - в этой версии реализации Я выбрал tf.nn.sigmoid_cross_entropy_with_logits: эта функция принимает двоичные данные и логиты параметров Бернулли в качестве аргументов, поэтому я не применил сигмовидную активацию к выходу последнего слоя декодера tfkl.Conv2DTranspose.

Обратите внимание, что эта функция сохраняет исходный размер входных данных: поскольку каждый экземпляр как цифровых изображений, так и выходных данных декодера имеет форму (28,28,1), tf.nn.sigmoid_cross_entropy_with_logit будет выводить тензор формы (batch_size, 28,28,1), где каждый элемент является отрицательной логарифмической вероятностью распределения Бернулли. для этого конкретного пикселя. Поскольку мы предполагаем, что каждый экземпляр цифрового изображения имеет независимое распределение Бернулли, отрицательная логарифмическая вероятность каждого экземпляра является суммой отрицательной логарифмической вероятности всех пикселей - отсюда и функция tf.math.reduce_sum(…, axis=[1, 2, 3]) в приведенном выше блоке кода.

Вот где следует принять дополнительные меры предосторожности, если вы планируете использовать Keras API: одно из самых больших преимуществ Keras API заключается в том, что он значительно упрощает код для обучения нейронной сети до трех строк: (1 ) построить объект tfk.Model, определив вход и выход сети; (2) compile модель, указав функцию loss и optimizer; (3) обучите модель, вызвав метод fit, аргументы которого включают входные и выходные данные, размер мини-пакета, количество эпох и т. д. Для этапа (2), аргумент loss принимает функцию ровно с двумя аргументами и вычисляет стоимость пакета данных во время обучения, взяв среднее значение по всем измерениям выходных данных этой функции, независимо от формы. Так что в нашем случае у вас может возникнуть соблазн написать loss=tf.nn.sigmoid_cross_entropy_with_logits на основе всего учебника по Keras, который вы видели, но это неверно, поскольку он берет среднее значение кросс-энтропийных потерь всех пикселей для каждого экземпляра, а не суммирует их. - итоговая стоимость больше не будет иметь статистической интерпретации. Но не бойтесь, вы все равно можете комбинировать tf.nn.sigmoid_cross_entropy_with_logits с Keras API - в версии c я подробно расскажу, как это сделать.

Все еще помните, в самом начале поста я упоминал, что в одном из собственных руководств TensorFlow была некорректная реализация VAE? А теперь пора присмотреться: первой ошибкой было вычисление ожидаемой ошибки реконструкции. Он применил mse_loss_fn = tfk.losses.MeanSquaredError() как функцию потерь: во-первых, выбор среднеквадратичной ошибки уже вызывает у меня сомнения - не только потому, что он неявно предполагает, что распределение реконструкции является независимым гауссовским с действительными данными, в то время как данные нашего изображения после предварительной обработки являются двоичный делает предположение о независимом распределении Бернулли более естественным выбором, но также и то, что MSE вычисляет масштабированную и сдвинутую отрицательную логарифмическую вероятность распределения Гаусса, что означает, что вы будете использовать потери Beta-VAE, не осознавая этого; и (гораздо) что более важно, tfk.losses.MeanSquaredError() без явного определения аргумента reduction также вычислит среднее значение MSE по всем измерениям. Опять же, единственное измерение, по которому нам нужно усреднить, - это измерение экземпляра, и для каждого экземпляра нам нужно взять сумму по всем пикселям. Если кто-то хочет применить модуль tfk.losses в реализации версии d., я продемонстрирую, как использовать tfk.losses.BinaryCrossentropy, более подходящий вариант для нашего случая для вычисления ожидаемой ошибки реконструкции.

  • Вычисление дивергенции KL

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

Обратите внимание, что для аналитического решения расхождения KL мы берем сумму параметров для апостериорного распределения z по их элементам по всем измерениям J. Однако в том же Руководстве по TensorFlow он вычисляется следующим образом:

kl_loss = -0.5 * tf.reduce_mean(
            z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1
        )

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

Можно также использовать приближение MC с одним образцом z для вычисления расхождения KL - проверьте код, когда analytic_kl=False. Если вы хотите проверить правильность своей реализации, вы можете использовать следующий способ аппроксимации расхождения KL с L, установленным как некоторое большое целое число, скажем 10 000:

в котором каждый z выбирается через

и посмотрите, похож ли результат на результат, полученный с помощью аналитического решения.

Наконец, если вы не хотите вручную вычислять расхождение KL, вы можете использовать функцию из библиотеки TensorFlow Probability, которая напрямую вычисляет расхождение KL между двумя распределениями:

prior_dist = tfd.MultivariateNormalDiag(loc=tf.zeros((batch_size, latent_dim)), scale_diag=tf.ones((batch_size, latent_dim)))
var_post_dist = tfd.MultivariateNormalDiag(loc=mu, scale_diag=sd)
kl_divergence = tfd.kl_divergence(distribution_a=var_post_dist, distribution_b=prior_dist)

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

Версия б.

Комбинируя TFP с Keras API из TF2, код выглядит намного проще, чем код в версии а. На самом деле это моя любимая версия, и я буду использовать ее в будущем из-за ее простоты.

В этой версии выходными данными и кодировщика, и декодера являются объекты из модуля tensorflow_probability.distributions, у которых есть много методов, которые можно ожидать от вероятностного распределения: mean, mode, sample, prob, log_prob и т. Д. Чтобы получить такой вывод от кодировщика и декодер, вам нужно только заменить их выходной слой одним из объектов tensorflow_probability.layers.

Реализация следующая:

Вот и все! 62 строки кода после переформатирования кода, включая комментарии. Одной из основных причин такого упрощения является выборка z, поскольку все этапы уловки повторной параметризации были абстрагированы через выходной уровень TFP кодировщика tfpl.IndependentNormal. Кроме того, вычисление дивергенции KL выполняется с помощью аргумента activity_regularizer в этом вероятностном выходном слое кодировщика, где мы указываем априорное распределение как стандартное многомерное гауссовское распределение, а также весовой коэффициент дивергенции KL ω, чтобы создать объект tfpl.KLDivergenceRegularizer. Кроме того, ожидаемую ошибку реконструкции можно вычислить, просто вызвав метод log_prob выходных данных декодера, который является объектом tfp.distributions.Independent - это настолько аккуратно, насколько это возможно.

Одно предостережение заключается в том, что можно подумать, поскольку вход в декодер должен быть тензором, но z является объектом tfp.distributions.Independent (см. Строку 53), нам нужно вместо этого написать z = encoder(x_input).sample() для явного примера z. Это не только не нужно, но и неправильно:

  • не нужно, потому что у нас convert_to_tensor_fn установлено в tfd.Distribution.sample (что на самом деле является значением по умолчанию, но я написал его явно, чтобы вы могли видеть): этот аргумент делает то, что всякий раз, когда вывод этого уровня обрабатывается как объект tf.Tensor, как в нашем случае когда нам понадобится образец из этого дистрибутива - outputs=decoder(z) в строке 56, он вызовет метод, указанный в convert_to_tensor_fn, поэтому он уже выполняет outputs=decoder(z.sample()).
  • неверно, потому что при явном вызове .sample() дивергенция KL, которая должна быть вычислена tfpl.KLDivergenceRegularizer, не будет включена в стоимость. Может случиться так, что после вызова .sample() у нас больше не будет z в качестве tfp.distributions.Independent объекта в вычислительном графе нейронной сети, который является типом объекта, содержащим tfpl.KLDivergenceRegularizer как его activity_regularizer. Следовательно, выполнение .sample() для вывода кодера в этой версии реализации VAE даст нам функцию потерь, которая содержит только ожидаемую ошибку реконструкции - обученный VAE больше не будет служить генеративной моделью, а будет только реконструировать свой ввод.

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

  • Они применили tfpl.MultivariateNormalTriL вместо tfpl.IndependentNormal в качестве выходного вероятностного слоя кодировщика, который по существу тренирует ненулевые элементы нижней треугольной матрицы, которая концептуально выводится из разложения Холецкого положительно определенной матрицы. Такая положительно определенная матрица по существу является ковариационной матрицей апостериорного распределения z и может быть любой положительно определенной матрицей, а не просто диагональной матрицей, принятой в статье VAE. Это дало бы нам более гибкое апостериорное распределение, но также содержало бы больше параметров для обучения, а расхождение KL сложнее вычислить.
  • Они установили расхождение KL weight по умолчанию 1.0 в tfpl.KLDivergenceRegularizer, но, как я обсуждал в части I, этот гиперпараметр имеет решающее значение для успеха реализации VAE и обычно требует быть явно настроенным для оптимизации производительности модели.

Наконец, для этой версии реализации я хочу показать одну особенность применения уровня TFP в качестве вывода декодера: то, что мы можем получать более гибкие прогнозы. Для исходного VAE выход декодера является детерминированным, поэтому после выборки z устанавливается выход декодера. Однако, когда вывод является распределением, мы можем вызвать метод mean, mode или sample для вывода объекта tf.Tensor в качестве предсказания цифрового изображения. Вот результаты, когда вызываются разные методы как для реконструкции, так и для генерации:

Обратите внимание, что как для реконструкции, так и для генерации среднее значение выглядит более размытым, чем режим (или не таким резким, как), поскольку среднее значение распределения Бернулли - это значение между 0 и 1. , в то время как режим равен 0 (если параметр меньше 0,5) или 1 (в противном случае); образец показывает большую степень детализации, чем два других, но также выглядит резким, поскольку каждый пиксель является случайным образцом (в отличие от режима, который после обучения в основном знает, какие значения будут принимать все пиксели в целом, чтобы сформировать цифру), который также принимает либо 0 или 1 в качестве его значения. Это гибкость, которую может обеспечить вероятностное распределение.

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

Версия c.

В версии a. мы говорили о том, как прямое использование tf.nn.sigmoid_cross_entropy_with_logits для loss аргумента, когда compile объект tfk.Model приведет к неправильному вычислению ожидаемой ошибки реконструкции; быстрое решение - реализовать пользовательскую функцию потерь на основе tf.nn.sigmoid_cross_entropy_with_logits следующим образом:

в котором мы берем сумму потерь кросс-энтропии по всем размерам пикселей для каждого экземпляра. А при составлении модели пишем

model.compile(loss=custom_sigmoid_cross_entropy_loss_with_logits, optimizer=tfk.optimizers.Adam(learning_rate))

Обратите внимание, что это тот же декодер, что и в версии a., который детерминированно выводит логиты параметров независимого распределения Бернулли. Между тем, мы используем tfpl.IndependentNormal в качестве выходного уровня кодировщика, как и в версии b., таким образом, вычисление расхождения KL выполняется с помощью его аргумента activity_regularizer.

Аналогично у нас есть следующая версия реализации:

Версия d.

Мы строим еще одну пользовательскую функцию потерь на tfk.losses.BinaryCrossentropy следующим образом:

Обратите внимание, что в отличие от версии c., в которой используется выходной уровень детерминированного декодера, в этой версии используется вероятностный уровень, как в версии b.; однако нам нужно сначала взять среднее значение этого распределения (строка 2 вышеприведенного блока кода), потому что одним из аргументов для объекта tfk.losses.BinaryCrossentropy является параметр распределения Бернулли, который совпадает с его иметь в виду. Также обратите внимание на аргумент reduction при инициализации объекта tfk.losses.BinaryCrossentropy, для которого установлено значение tfk.losses.Reduction.NONE: это предотвращает выполнение программой дальнейших операций с результирующим тензором, имеющим ту же форму, что и тензор изображения мини-пакета, каждый элемент которого содержит потеря кросс-энтропии для одного конкретного пикселя. Затем мы берем сумму по размерам на уровне экземпляра, как и в пользовательской функции потерь в версии c.

Стоит отметить, что любая функция, которую мы определяем для loss аргумента, когда compile модель, должна принимать ровно два аргумента, один из которых является данными, которые модель пытается предсказать, а другой - выходными данными модели. Следовательно, у нас возникнет проблема, когда мы захотим применить Keras API, имея более гибкую функцию потерь. Нам повезло иметь аргумент activity_regularizer в выходном уровне TFP кодировщика, который помогает учесть расхождение KL при вычислении стоимости, что дает нам версию b.; но что, если мы этого не сделаем? В оставшихся двух версиях представлен способ реализации более гибкой функции потерь в целом, а не только для конкретного случая VAE - благодаря методу add_loss класса tfkl.Layer.

Версия e.

Я начну с демонстрации кода для этой версии:

Обратите внимание, что я использовал ту же функцию loss, что и в версии c., когда compile tfk.Model (строка 70). Взвешенное расхождение KL учитывается при вычислении стоимости путем вызова метода класса add_loss (строка 49). Ниже приводится прямая цитата из его документации:

Этот метод можно использовать внутри подкласса слоя или функции call модели, и в этом случае losses должен быть тензором или списком тензоров.

Без каких-либо требований к формату, кроме «тензора или списка тензоров», мы можем построить гораздо более гибкую функцию потерь. Сложная часть проистекает из первой половины этой цитаты, в которой указывается, где должен быть вызван этот метод: обратите внимание, что в моей реализации я вызвал его внутри функции call класса VAE_MNIST, который наследуется от класса tfk.Model. Это отличается от реализации версии b, в которой нет наследования классов. Для этой версии я изначально также начал с написания класса VAE_MNIST без наследования классов и вызова метода compile для объекта tfk.Model в функции класса с именем build_vae_keras_model, точно так же, как то, что я сделал в версии b. tfk.Model объект model компилируется, я напрямую вызвал model.add_loss(self.kl_weight * kl_divergence) - мое объяснение состоит в том, что, поскольку выходной слой нашей модели, объект класса tfkl.Conv2DTranspose, наследуется от tfkl.Layer, мы сможем провести такую ​​операцию. Однако в результате тренировки я получу «более яркое» изображение, если уменьшу вес KL примерно до 0,01; и если я начну увеличивать вес KL, я получу почти полностью темные сгенерированные изображения. В целом сгенерированные изображения будут иметь низкое качество, в то время как восстановленные изображения будут выглядеть великолепно, но никак не изменится при разных значениях веса KL. Только после того, как я унаследовал VAE_MNIST от класса tfk.Model и изменил имя метода класса, в котором вызывается add_loss, на call, модель будет правильно обучаться.

Версия f.

Я реализовал эту версию как последнюю для практики, полагая, что такая структура может пригодиться, если однажды мне понадобится реализовать какую-то модель, которая слишком сложна для использования Keras API. Вот код:

Я хочу подчеркнуть важность наименования метода класса, который вызывает метод add_loss, как call: обратите внимание, что этот метод изначально был назван encode_and_decode (Line 70) - если я добавлю @tf.function декоратор для метода train_step (Line 94) перед переименованием метода класса я бы получил

TypeError: An op outside of the function building code is being passed a “Graph” tensor. It is possible to have Graph tensors leak out of the function building context by including a tf.init_scope in your function building code.

После того как я соответствующим образом изменил имя метода класса, ошибка исчезла. Важно иметь декоратор @tf.function, потому что он значительно увеличивает скорость обучения: перед переименованием последняя эпоха длилась 1041 секунд (и промежуток времени для каждой эпохи увеличивался от эпохи к эпохе; первая эпоха « всего »длилось 139 секунд, что является еще одним странным явлением); после переименования каждая эпоха имеет время выполнения от 34,7 секунды до 39,4 секунды.

Более того, перед переименованием мне нужно было бы добавить аргумент lambda для метода add_loss (строка 78), поскольку без lambda: я бы получил следующую ошибку:

ValueError: Expected a symbolic Tensors or a callable for the loss value. Please wrap your loss computation in a zero argument lambda.

Короче говоря, когда вы планируете применить метод add_loss для реализации более гибкой функции потерь, рекомендуется создать подкласс класса tfk.Model и вызвать метод add_loss внутри метода класса с именем call.

Заключение

В этом посте я представил 6 различных способов реализации модели VAE, которая обучается на наборе данных MNIST, и дал подробный обзор VAE с практической точки зрения. В общем, я бы порекомендовал начинающему внедрять VAE начать с версии a., а затем попытаться применить Keras API для упрощения кода с настраиваемой функцией потерь для вычисления ожидаемой ошибки реконструкции, как и Версия c. и Версия d.; После того, как вы получите хорошее представление о модели, было бы неплохо принять версию b., которая является наиболее краткой и, IMHO, самой элегантной версией. Если кто-то хочет подготовиться к реализации очень гибкой функции потерь в целом в будущем, можно начать с версии e. или версии f..

использованная литература

[1] Дидерик П. Кингма и Макс Веллинг, Автоматическое кодирование вариационного байесовского протокола (2013), Труды 2 Международной конференции по обучающим представлениям (ICLR)

[2] Чарльз Бланделл, Жюльен Корнебис, Корай Кавукчуоглу и Даан Виерстра, Неопределенность веса в нейронных сетях (2015), Труды 32 Международной конференции по машинному обучению (ICML)

[3] Создание новых слоев и моделей посредством создания подклассов (2020), Руководство по TensorFlow

[4] Сверточный вариационный автоэнкодер (2020), Учебное пособие по TensorFlow

[5] Ян Фишер, Алекс Алеми, Джошуа В. Диллон и команда TFP, Вариационные автоэнкодеры с вероятностными слоями Tensorflow (2019), TensorFlow on Medium