Anycost GAN

视频 | 论文 | 网站

麻省理工学院，Adobe研究院，卡内基梅隆大学

发表于CVPR 2021

灵活性

Anycost GAN在不同计算预算下生成一致的输出。

演示

在这里，我们可以使用Anycost生成器进行交互式图像编辑。完整的生成器需要约3秒来渲染一张图像，这对于编辑来说太慢了。而使用Anycost生成器，我们可以以5倍的速度提供视觉上相似的预览。调整后，我们点击"完成"按钮来合成高质量的最终输出。点击这里查看完整演示。

概述

Anycost生成器可以通过使用不同的通道和分辨率配置，以多样化的计算成本运行。子生成器与完整生成器相比，能够实现高度的输出一致性，提供快速预览。

通过(1)基于采样的多分辨率训练，(2)自适应通道训练，以及(3)生成器条件判别器，我们在不同分辨率和通道下实现了高图像质量和一致性。

结果

Anycost GAN（均匀通道版本）支持4种分辨率和4种通道比率，产生视觉上一致但图像保真度不同的图像。

在图像投影和编辑过程中，这种一致性得以保持：

使用方法

入门

克隆此仓库：

git clone https://github.com/mit-han-lab/anycost-gan.git
cd anycost-gan

安装PyTorch 1.7和其他依赖项。

我们建议使用Anaconda设置环境：conda env create -f environment.yml

介绍性笔记本

我们提供了一个Jupyter笔记本示例，展示如何使用Anycost生成器以多样化成本进行图像合成：notebooks/intro.ipynb。

我们还提供了笔记本的Colab版本：。请确保在运行时选项中选择GPU作为加速器。

交互式演示

我们提供了一个交互式演示，展示如何使用Anycost GAN实现交互式图像编辑。运行演示：

python demo.py

如果您的计算机包含CUDA GPU，请尝试运行：

FORCE_NATIVE=1 python demo.py

您可以在这里找到演示的视频录像。

使用预训练模型

要获取预训练的生成器、编码器和编辑方向，请运行：

import models

pretrained_type = 'generator'  # 从['generator', 'encoder', 'boundary']中选择
config_name = 'anycost-ffhq-config-f'  # 替换为其他模型的配置名称
models.get_pretrained(pretrained_type, config=config_name)

我们还提供了面部属性分类器（适用于不同的生成器）用于计算编辑方向。您可以通过运行以下命令获取：

models.get_pretrained('attribute-predictor')

属性分类器接受FFHQ格式的面部图像作为输入。

加载Anycost生成器后，我们可以在广泛的计算成本范围内运行它。例如：

from models.dynamic_channel import set_uniform_channel_ratio, reset_generator

g = models.get_pretrained('generator', config='anycost-ffhq-config-f')  # anycost uniform
set_uniform_channel_ratio(g, 0.5)  # 设置通道
g.target_res = 512  # 设置分辨率
out, _ = g(...)  # 生成图像
reset_generator(g)  # 恢复生成器

有关详细用法和灵活通道Anycost生成器，请参阅notebooks/intro.ipynb。

模型库

目前，我们提供以下预训练的生成器、编码器和编辑方向。我们将在未来添加更多。

对于Anycost生成器，默认情况下，我们指的是均匀设置。

配置名称	生成器	编码器	编辑方向
anycost-ffhq-config-f	:heavy_check_mark:	:heavy_check_mark:	:heavy_check_mark:
anycost-ffhq-config-f-flexible	:heavy_check_mark:	:heavy_check_mark:	:heavy_check_mark:
anycost-car-config-f	:heavy_check_mark:
stylegan2-ffhq-config-f	:heavy_check_mark:	:heavy_check_mark:	:heavy_check_mark:

stylegan2-ffhq-config-f指的是从仓库转换的官方StyleGAN2生成器。

数据集

我们将FFHQ、CelebA-HQ和LSUN Car数据集准备成一个图像目录，以便可以轻松地与torchvision中的ImageFolder一起使用。数据集布局如下：

├── PATH_TO_DATASET
│   ├── images
│   │   ├── 00000.png
│   │   ├── 00001.png
│   │   ├── ...

由于版权问题，您需要从官方网站下载数据集并进行相应处理。

评估

我们提供了评估论文中提到的一些指标的代码。部分代码使用horovod编写，以支持分布式评估并减少GPU间通信成本，大大提高了速度。请查看其网站以进行正确安装。

Fre ́chet Inception Distance (FID)

在评估FID之前，您需要使用以下脚本计算真实图像的inception特征：

python tools/calc_inception.py \
    --resolution 1024 --batch_size 64 -j 16 --n_sample 50000 \
    --save_name assets/inceptions/inception_ffhq_res1024_50k.pkl \
    PATH_TO_FFHQ

或者您可以从这里下载预先计算好的inception特征，并将其放在assets/inceptions下。

然后，您可以通过运行以下命令来评估FID：

horovodrun -np N_GPU \
    python metrics/fid.py \
    --config anycost-ffhq-config-f \
    --batch_size 16 --n_sample 50000 \
    --inception assets/inceptions/inception_ffhq_res1024_50k.pkl
    # --channel_ratio 0.5 --target_res 512  # 可选择使用较小的分辨率/通道数

Perceptual Path Length (PPL)

类似地，通过以下命令评估PPL：

horovodrun -np N_GPU \
    python metrics/ppl.py \
    --config anycost-ffhq-config-f

属性一致性

通过运行以下命令评估属性一致性：

horovodrun -np N_GPU \
    python metrics/attribute_consistency.py \
    --config anycost-ffhq-config-f \
    --channel_ratio 0.5 --target_res 512  # 子生成器的配置；必需

编码器评估

要评估编码器的性能，请运行：

python metrics/eval_encoder.py \
    --config anycost-ffhq-config-f \
    --data_path PATH_TO_CELEBA_HQ

训练

我们提供了在FFHQ数据集上训练Anycost GAN的脚本。

在FFHQ上训练原始StyleGAN2

horovodrun -np 8 bash scripts/train_stylegan2_ffhq.sh

原始StyleGAN2的训练非常耗时。我们建议从这里下载转换后的检查点，并将其放在checkpoint/下。

训练Anycost GAN：多分辨率

horovodrun -np 8 bash scripts/train_stylegan2_multires_ffhq.sh

请注意，在每个epoch之后，我们会评估两种分辨率（1024和512）的FID，以更好地监控训练进度。我们还应用了蒸馏技术来加速收敛，这在论文的消融实验中没有使用。

训练Anycost GAN：自适应通道

horovodrun -np 8 bash scripts/train_stylegan2_multires_adach_ffhq.sh

这里我们设置了更长的训练周期以获得更稳定的复现效果，这可能并非必要（取决于随机性）。

注意：我们在具有24GB内存的Titan RTX GPU上训练了我们的模型。对于内存较小的GPU，您可能需要降低分辨率/模型大小/批量大小等，并相应地调整其他超参数。

引用

如果您在研究中使用了此代码，请引用我们的论文。

@inproceedings{lin2021anycost,
  author    = {Lin, Ji and Zhang, Richard and Ganz, Frieder and Han, Song and Zhu, Jun-Yan},
  title     = {Anycost GANs for Interactive Image Synthesis and Editing},
  booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year      = {2021},
}