斯坦尼斯拉夫·皮德霍尔斯基 • 唐纳德·A·阿杰罗 • 吉安弗兰科·多雷托
包含模型和定性结果的 Google Drive 文件夹
对抗性潜在自动编码器
斯坦尼斯拉夫·皮德霍斯基、唐纳德·阿杰罗、吉安弗兰科·多雷托摘要:自动编码器网络是一种无监督方法,旨在通过同时学习编码器-生成器映射来结合生成属性和表示属性。尽管进行了广泛的研究,但它们是否具有与 GAN 相同的生成能力或学习解缠结表示的问题尚未得到充分解决。我们引入了一种联合解决这些问题的自动编码器,我们称之为对抗性潜在自动编码器(ALAE)。它是一个通用架构,可以利用 GAN 训练程序的最新改进。我们设计了两种自动编码器:一种基于 MLP 编码器,另一种基于 StyleGAN 生成器,我们称之为 StyleALAE。我们验证了两种架构的解缠结特性。我们证明,StyleALAE 不仅可以生成与 StyleGAN 质量相当的 1024x1024 人脸图像,而且在相同分辨率下还可以生成基于真实图像的人脸重建和操作。这使得 ALAE 成为第一个能够与仅生成器类型的架构进行比较并超越其功能的自动编码器。
@InProceedings{pidhorskyi2020adversarial,
author = {Pidhorskyi, Stanislav and Adjeroh, Donald A and Doretto, Gianfranco},
booktitle = {Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR)},
title = {Adversarial Latent Autoencoders},
year = {2020},
note = {[to appear]},
}
要运行演示,您需要安装支持 CUDA 的 GPU、PyTorch >= v1.3.1 和 cuda/cuDNN 驱动程序。安装所需的软件包:
pip install -r requirements.txt
下载预训练模型:
python training_artifacts/download_all.py
运行演示:
python interactive_demo.py
您可以指定要使用的yaml配置。配置位于:https://github.com/podgorskiy/ALAE/tree/master/configs。默认情况下,它使用 FFHQ 数据集。您可以使用-c
参数更改配置。要以 256x256 分辨率在celeb-hq
上运行,请运行:
python interactive_demo.py -c celeba-hq256
但是,对于 FFHQ 以外的配置,您需要获取属性的新主方向向量。
存储库中的代码的组织方式是所有脚本都必须从存储库的根目录运行。如果您使用IDE(例如PyCharm或Visual Studio Code),只需将工作目录设置为指向存储库的根目录。
如果你想从命令行运行,那么你还需要设置PYTHONPATH变量以指向存储库的根目录。
例如,假设我们已将存储库克隆到~/ALAE目录,然后执行以下操作:
$ cd ~/ALAE
$ export PYTHONPATH=$PYTHONPATH:$(pwd)
现在您可以按如下方式运行脚本:
$ python style_mixing/stylemix.py
小路 | 描述 |
---|---|
ALAE | 存储库根文件夹 |
├ 配置 | 包含 yaml 配置文件的文件夹。 |
│ ├ 卧室.yaml | 分辨率为 256x256 的 LSUN 卧室数据集的配置文件。 |
│ ├ celeba.yaml | 128x128 分辨率的 CelebA 数据集的配置文件。 |
│ ├ celeba-hq256.yaml | 256x256 分辨率的 CelebA-HQ 数据集的配置文件。 |
│ ├ celeba_ablation_nostyle.yaml | 用于消融研究的 CelebA 128x128 数据集的配置文件(无样式)。 |
│ ├ celeba_ablation_separate.yaml | 用于消融研究的 CelebA 128x128 数据集的配置文件(单独的编码器和鉴别器)。 |
│ ├ celeba_ablation_z_reg.yaml | 用于消融研究的 CelebA 128x128 数据集的配置文件(Z 空间回归,而不是 W 空间)。 |
│ ├ ffhq.yaml | 1024x1024 分辨率的 FFHQ 数据集的配置文件。 |
│ ├ mnist.yaml | 使用 Style 架构的 MNIST 数据集的配置文件。 |
│ └ mnist_fc.yaml | 仅使用全连接层(排列不变 MNIST)的 MNIST 数据集的配置文件。 |
├ 数据集_准备 | 包含用于数据集准备的脚本的文件夹。 |
│ ├ 准备_celeba_hq_tfrec.py | 为 CelebA-HQ 数据集准备 256x256 分辨率的 TFRecords。 |
│ ├ 准备_celeba_tfrec.py | 为 CelebA 数据集准备 128x128 分辨率的 TFRecords。 |
│ ├ 准备_mnist_tfrec.py | 为 MNIST 数据集准备 TFRecords。 |
│ ├ split_tfrecords_bedroom.py | 将官方 TFRecords 从 StyleGAN 论文中拆分为 LSUN 卧室数据集。 |
│ └ split_tfrecords_ffhq.py | 将官方 TFRecords 从 StyleGAN 论文中拆分为 FFHQ 数据集。 |
├ 数据集样本 | 包含不同数据集的示例输入的文件夹。用于训练期间的数字和测试输入。 |
├ make_figures | 用于制作各种人物的脚本。 |
├ 指标 | 用于计算指标的脚本。 |
├ 主体方向 | 用于计算各种属性的主方向向量的脚本。用于交互式演示。 |
├ 风格_混合 | 用于生成风格混合图形的示例输入和脚本。 |
├ 训练_工件 | 保存检查点/样本输出/绘图的默认位置。 |
│ └ download_all.py | 用于下载所有预训练模型的脚本。 |
├ 交互演示.py | 用于交互式演示的可运行脚本。 |
├ train_alae.py | 用于训练的可运行脚本。 |
├ train_alae_separate.py | 用于训练消融研究的可运行脚本(单独的编码器和鉴别器)。 |
├ 检查点.py | 用于保存/恢复模型权重、优化器状态和损失历史记录的模块。 |
├ 自定义_adam.py | 定制的 Adam 优化器,用于学习率均衡和零秒 Beta。 |
├ 数据加载器.py | 具有数据集类、加载器、迭代器等的模块。 |
├ 默认值.py | 具有默认值的配置变量的定义。 |
├ 启动器.py | 用于运行多 GPU、多进程训练的助手。设置配置和日志记录。 |
├ lod_driver.py | 用于管理增长/稳定网络的助手类。 |
├ lreq.py | 用于学习率均衡的自定义Linear 、 Conv2d 和ConvTranspose2d 模块。 |
├ 模型.py | 具有高级模型定义的模块。 |
├ model_separate.py | 与上面相同,但用于消融研究。 |
├ 网络.py | 多种架构的所有网络块的定义。 |
├ 注册表.py | 用于从配置文件中选择的网络块注册表。 |
├ 调度程序.py | 具有热启动和聚合多个优化器的自定义调度程序。 |
├ 追踪器.py | 用于绘制损失的模块。 |
└ 实用程序.py | 用于异步调用的装饰器、用于缓存的装饰器、用于网络块的注册表。 |
在此代码库中, yacs用于处理配置。
大多数可运行脚本接受-c
参数,该参数可以指定要使用的配置文件。例如,要制作重建图形,您可以运行:
python make_figures/make_recon_figure_paged.py
python make_figures/make_recon_figure_paged.py -c celeba
python make_figures/make_recon_figure_paged.py -c celeba-hq256
python make_figures/make_recon_figure_paged.py -c bedroom
默认配置是ffhq
。
训练是使用 TFRecords 完成的。 TFRecords 使用 DareBlopy 读取,它允许将它们与 Pytorch 一起使用。
在配置文件以及所有准备脚本中,假设所有数据集都位于/data/datasets/
中。您可以更改配置文件中的路径,也可以创建指向存储数据集的位置的符号链接。
生成 CelebA-HQ 的官方方式可能具有挑战性。请参考此页面:https://github.com/suvojit-0x55aa/celebA-HQ-dataset-download 您可以从以下位置获取预先生成的数据集:https://drive.google.com/drive/folders/11Vz0fqHS2rXDb5pprgTjpD7S2BAJhi1P
要下载预训练模型,请运行:
python training_artifacts/download_all.py
注意:由于下载限制,过去从 Google Drive 下载模型时会出现问题。现在,脚本的设置方式是,如果无法从 Google Drive 下载数据,它将尝试从 S3 下载数据。
如果遇到问题,请尝试删除所有 *.pth 文件,更新dlutils软件包( pip install dlutils --upgrade
),然后再次运行download_all.py
。如果这不能解决问题,请提出问题。另外,您可以尝试从此处手动下载模型:https://drive.google.com/drive/folders/1tsI1q1u8QRX5t7_lWCSjpniLGlNY-3VY?usp=sharing
在配置文件中, OUTPUT_DIR
指向保存和读取权重的位置。例如: OUTPUT_DIR: training_artifacts/celeba-hq256
在OUTPUT_DIR
中,它保存一个文件last_checkpoint
,其中包含带有模型权重的实际.pth
pickle 的路径。如果你想用特定的权重文件测试模型,你可以简单地修改last_checkpoint
文件。
要生成风格混合图形,请运行:
python style_mixing/stylemix.py -c <config>
其中代替<config>
的是: ffhq
、 celeba
、 celeba-hq256
、 bedroom
之一
要使用多尺度图像生成重建:
python make_figures/make_recon_figure_multires.py -c <config>
要从多个页面上的所有样本输入生成重建:
python make_figures/make_recon_figure_paged.py -c <config>
还有:
python make_figures/old/make_recon_figure_celeba.py
python make_figures/old/make_recon_figure_bed.py
要从 FFHQ 测试集生成重建:
python make_figures/make_recon_figure_ffhq_real.py
要生成插值图:
python make_figures/make_recon_figure_interpolation.py -c <config>
生成遍历图:
(对于 FFHQ 以外的数据集,您需要首先找到主要方向)
python make_figures/make_traversarls.py -c <config>
要使生成图运行:
make_generation_figure.py -c <config>
除了安装所需的包之外:
pip install -r requirements.txt
您需要安装 DareBlopy:
pip install dareblopy
运行训练:
python train_alae.py -c <config>
它将在所有可用的 GPU 上运行多 GPU 训练。它使用DistributedDataParallel
实现并行性。如果只有一个 GPU 可用,它将在单个 GPU 上运行,无需特别注意。
建议的 GPU 数量为 8 个。较少数量的 GPU 上的再现性可能会出现问题。您可能需要根据 GPU 的内存大小调整配置文件中的批处理大小。
除了安装所需的软件包和 DareBlopy 之外,您还需要从 StyleGAN 安装 TensorFlow 和 dnnlib。
Tensorflow 版本必须是1.10
:
pip install tensorflow-gpu==1.10
它需要 CUDA 版本 9.0。
也许,最好的方法是使用 Anaconda 来处理这个问题,但我更喜欢从 pop-os 存储库安装 CUDA 9.0(适用于 Ubuntu):
sudo echo "deb http://apt.pop-os.org/proprietary bionic main" | sudo tee -a /etc/apt/sources.list.d/pop-proprietary.list
sudo apt-key adv --keyserver keyserver.ubuntu.com --recv-key 204DD8AEC33A7AFF
sudo apt update
sudo apt install system76-cuda-9.0
sudo apt install system76-cudnn-9.0
然后只需设置LD_LIBRARY_PATH
变量:
export LD_LIBRARY_PATH=/usr/lib/cuda-9.0/lib64
Dnnlib 是 StyleGAN 中使用的包。您可以使用以下命令安装它:
pip install https://github.com/podgorskiy/dnnlib/releases/download/0.0.1/dnnlib-0.0.1-py3-none-any.whl
用于运行指标的所有代码很大程度上基于 StyleGAN 存储库中的代码。它还使用相同的预训练模型:
https://github.com/NVlabs/stylegan#licenses
inception_v3_features.pkl 和 inception_v3_softmax.pkl 源自 Christian Szegedy、Vincent Vanhoucke、Sergey Ioffe、Jonathon Shlens 和 Zbigniew Wojna 预训练的 Inception-v3 网络。该网络最初是在 TensorFlow 模型存储库上根据 Apache 2.0 许可证共享的。
vgg16.pkl 和 vgg16_zhang_perceptual.pkl 源自 Karen Simonyan 和 Andrew Zisserman 预训练的 VGG-16 网络。该网络最初是在用于大规模视觉识别的非常深卷积网络项目页面上根据 Creative Commons BY 4.0 许可证共享的。
vgg16_zhang_perceptual.pkl 进一步源自 Richard Zhu、Phillip Isola、Alexei A. Efros、Eli Shechtman 和 Oliver Wang 预训练的 LPIPS 权重。这些权重最初是在 PerceptualSimilarity 存储库上的 BSD 2-Clause“简化”许可证下共享的。
最后,运行指标:
python metrics/fid.py -c <config> # FID score on generations
python metrics/fid_rec.py -c <config> # FID score on reconstructions
python metrics/ppl.py -c <config> # PPL score on generations
python metrics/lpips.py -c <config> # LPIPS score of reconstructions