Два этапа VAE

Это код модели двухэтапного VAE, предложенной в нашей статье ICLR 2019 «Диагностика и улучшение моделей VAE» [1].

[1] Dai, B. and Wipf, D. Diagnosing and enhancing VAE models. In International Conference on Learning Representations, 2019.

Шаг 1. Подготовка набора данных

Мы проводим эксперименты на наборах данных MNIST, Fashion-MNIST, Cifar-10 и CelebA. Эти данные загружаются с официального веб-сайта, а затем преобразуются в формат npy с помощью скрипта preprocess.py. Следуйте следующим шагам для подготовки каждого набора данных. Или вы можете напрямую загрузить данные из Google Doc (если вы загружаете данные непосредственно из Google Doc, извлеките файл в корневую папку).

MNIST

Загрузите данные с: http://yann.lecun.com/exdb/mnist/.

Вы получите файлы t10k-images-idx3-ubyte, t10k-labels-idx1-ubyte, train-images-idx3-ubyte и train-labels-idx1-ubyte. Поместите их в папку ./data/mnist.

Fashion-MNIST

Скачайте данные с: https://github.com/zalandoresearch/fashion-mnist.

Снова вы получите четыре файла t10k-images-idx3-ubyte, t10k-labels-idx1-ubyte, train-images-idx3-ubyte и train-labels-idx1-ubyte. Положите их в папку ./data/fashion.

Для предварительной обработки MNIST и Fashion-MNIST вам также необходимо установить пакет python-mnist с помощью команды:

pip install python-mnist

Cifar-10

Скачать данные (версия Python) с: https://www.cs.toronto.edu/~kriz/cifar.html.

Извлеките скачанный файл в ./data/cifar10. Там будет одна подпапка под названием cifar-10-batches-py. Внутри этой папки будет 6 файлов с именами data_batch_1, data_batch_2, data_batch_3, data_batch_4, data_batch_5 и test_batch.

Набор данных CelebA

Скачивайте данные с: http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html.

Поместите извлечённые файлы изображений (202 599 файлов jpg) в папку ./data/celeba/img_align_celeba.

Предварительная обработка

Преобразуйте данные в формат .npy, выполнив команду:

python preprocess.py

Вы получите несколько файлов .npy в каждой подпапке под ./data.

Шаг 2. Запуск двухэтапного VAE

Обучите модель, запустив:

python demo.py --dataset [DATASET] --network-structure [NETWORK] --exp-name [EXP] --gpu [GPU] Аргумент --exp-name позволяет вам настроить идентификатор вашего эксперимента. Он создаст папку ./experiments/[DATASET]/[EXP] для размещения всех выходных файлов и изображений. Если вы хотите только протестировать модель, добавьте --val в командную строку.

• Наборы данных: используйте строчные буквы. Вы можете выбрать из: mnist, fashion, cifar10 и celeba.
• Структура сети: вы можете выбрать из Infogan, Wae и Resnet. Чтобы реализовать свои собственные сетевые структуры, вы можете добавить новый производный класс TwoStageVaeModel в ./network/two_stage_vae_model.py и просто реализовать функции build_encoder1 и build_decoder1.
• Потеря кросс-энтропии: если вы хотите провести эксперименты с потерей кросс-энтропии (метод VAE(cross-entr.) в статье), вы можете добавить --cross-entropy-loss в командной строке.

Сгенерированные образцы

Чтобы воспроизвести следующие результаты с архитектурой Resnet, запустите:

python demo.py --dataset celeba --epochs 100 --lr-epochs 40 --epochs2 100 --lr-epochs2 40 --network-structure Resnet --num-scale 4 --base-dim 32 --latent-dim 128 --gpu [GPU] --exp-name [EXP]

• Случайные выборки, сгенерированные первым этапом VAE с использованием структуры Resnet на CelebA:

• Случайные выборки, сгенерированные вторым этапом VAE с использованием структуры Resnet на CelebA:

Чтобы воспроизвести следующие результаты с сетевой структурой WAE, запустите:

python dome.py --dataset celeba --epochs 70 --lr-epochs 30 --epochs2 70 --lr-epochs2 30 --network-structure Wae --gpu [GPU] --exp-name [EXP]

• Случайные выборки, сгенерированные первым этапом VAE с использованием сетевой структуры WAE на CelebA:

• Случайные выборки, сгенерированные вторым этапом VAE с использованием сетевой структуры WAE на CelebA:

stage](imgs/wae_gen2.jpg)

сцена (imgs/wae gen2.jpg)