HyperInverter: 通过超网络改进StyleGAN反演

Tan M. Dinh, Anh Tran, Rang Nguyen, Binh-Son Hua<br> VinAI Research, 越南

摘要: 近年来,得益于对GAN潜在空间的探索和利用,真实世界图像操作取得了惊人的进展。GAN反演是这一流程的第一步,旨在忠实地将真实图像映射到潜在代码。不幸的是,大多数现有的GAN反演方法无法同时满足以下三个要求中的至少一个:高重建质量、可编辑性和快速推理。在这项研究中,我们提出了一种新颖的两阶段策略,可以同时满足所有要求。在第一阶段,我们训练一个编码器将输入图像映射到StyleGAN2的W空间,该空间被证明具有出色的可编辑性但较低的重建质量。在第二阶段,我们利用一系列超网络来恢复反演过程中丢失的信息,从而补充初始阶段的重建能力。这两个步骤相互补充,由于超网络分支而产生高重建质量,并由于W空间中完成的反演而具有出色的可编辑性。我们的方法完全基于编码器,从而实现极快的推理。在两个具有挑战性的数据集上进行的大量实验证明了我们方法的优越性。


我们的方法在质量上显著优于其他基于编码器的方法(pSp、e4e、ReStyle),同时推理时间与它们相同。与基于优化的方法(SG2-W+、PTI)相比,我们的工作在质量上与SG2-W+相当,略低于PTI。然而,值得注意的是,我们的方法运行速度非常快,可以用于交互式应用(比SG2-W+和PTI分别快3000倍和1100倍)。

模型架构和实验结果的详细信息可以在我们的论文中找到。

@inproceedings{dinh2021hyperinverter,
    title={HyperInverter: Improving StyleGAN Inversion via Hypernetwork},
    author={Tan M. Dinh and Anh Tuan Tran and Rang Nguyen and Binh-Son Hua},
    booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year={2022}
}

请在使用我们的模型实现来产生发表结果或将其纳入其他软件时引用我们的论文。

开始使用

代码库在以下环境中测试:

Ubuntu
CUDA 10.0, CuDNN 7

安装

克隆此仓库:

git clone https://github.com/VinAIResearch/HyperInverter.git
cd HyperInverter

安装依赖:

conda create -p ./envs python=3.7.3
conda activate ./envs
pip install -r requirements.txt

数据集

人脸: 我们使用FFHQ数据集中的70,000张图像进行训练,使用CelebA-HQ数据集中的2,824张图像进行测试。图像分辨率为1024 x 1024,并裁剪对齐到中心。有关预处理步骤的更多详细信息,请参考FFHQ。
教堂: 我们使用LSUN Church官方训练集和测试集中的所有126,227张图像和300张图像来训练和评估我们的模型。图像调整为256 x 256分辨率。

请下载相应的数据集并解压到data文件夹。然后,转到configs/paths_config.py并修改此文件以正确链接数据:

dataset_paths = {
    "ffhq": "/path/to/ffhq/train_img",
    "celeba_test": "/path/to/CelebA-HQ/test_img",
    "church_train": "/path/to/lsun-church/train_img",
    "church_test": "/path/to/lsun-church/test_img",
}

如果您需要尝试自己的数据集,可以在以下文件中进行必要的修改:(i) data_configs.py定义数据路径;(ii) transforms_configs.py定义数据转换。

辅助预训练模型

运行以下命令自动下载实验所需的辅助预训练模型。

python scripts/download_auxiliary_pretrained_models.py

基本上,我们将这些路径正确设置为训练或推理过程。如果您想更改这些路径,请访问文件configs/paths_config.py并修改model_paths字典中的相应记录。

实验

预训练模型

有关我们官方预训练模型的信息,请参阅模型库。请下载我们的预训练模型(W编码器和HyperInverter)并将它们放入pretrained_models文件夹。

训练

第一阶段: W编码器

请按照pSp或e4e的说明训练W编码器,将图像编码为W空间(512)中的潜在代码。

为了节省时间,我们在模型库中发布了用于教堂和人脸的预训练W编码器。

第二阶段: HyperInverter

我们提供了默认的训练脚本,如下所示。如果您有时间,请进一步调整超参数以获得最佳结果。请注意,参数--hidden_dim对模型质量有很大影响,增加它会导致更大的模型,并倾向于有更好的质量。在论文中,我们使用hidden_dim=256。但是,我们建议使用hidden_dim=128来平衡模型大小和模型性能。此外,值得注意的是,我们的代码也很容易修改以预测StyleGAN2的不同层权重。参考weight_shapes并根据需要进行修改。然后,修改hyper_inverter中的get_target_shapes函数以添加您的选项。

人脸

EXPERIMENT_DIR=""
W_ENCODER_PATH=""
GPU_ID=0
CUDA_VISIBLE_DEVICES="$GPU_ID" \
python scripts/train.py \
--dataset_type=ffhq_encode \
--encoder_type=LayerWiseEncoder \
--w_encoder_path="$W_ENCODER_PATH" \
--output_size=1024 \
--exp_dir="$EXPERIMENT_DIR" \
--batch_size=8 \
--batch_size_used_with_adv_loss=4 \
--workers=4 \
--val_interval=1000 \
--save_interval=5000 \
--encoder_optim_name=adam \
--discriminator_optim_name=adam \
--encoder_learning_rate=1e-4 \
--discriminator_learning_rate=1e-4 \
--hyper_lpips_lambda=0.8 \
--hyper_l2_lambda=1.0 \
--hyper_id_lambda=0.1 \
--hyper_adv_lambda=0.005 \
--hyper_d_reg_every=16 \
--hyper_d_r1_gamma=10.0 \
--step_to_add_adversarial_loss=200000 \
--target_shape_name=conv_without_bias  \
--max_steps=500000 \
--hidden_dim=128 \
--num_cold_steps=20000 \
--save_checkpoint_for_resuming_training \
--use_wandb

教堂

EXPERIMENT_DIR=""
W_ENCODER_PATH=""
GPU_ID=0

CUDA_VISIBLE_DEVICES="$GPU_ID" \
python scripts/train.py \
--dataset_type=church_encode \
--encoder_type=ResNetLayerWiseEncoder \
--w_encoder_path="$W_ENCODER_PATH" \
--output_size=256 \
--exp_dir="$EXPERIMENT_DIR" \
--batch_size=8 \
--batch_size_used_with_adv_loss=4 \
--workers=4 \
--val_interval=1000 \
--save_interval=5000 \
--encoder_optim_name=adam \
--discriminator_optim_name=adam \
--encoder_learning_rate=1e-4 \
--discriminator_learning_rate=1e-4 \
--hyper_lpips_lambda=0.8 \
--hyper_l2_lambda=1.0 \
--hyper_id_lambda=0.5 \
--hyper_adv_lambda=0.15 \
--hyper_d_reg_every=16 \
--hyper_d_r1_gamma=100.0 \
--step_to_add_adversarial_loss=100000 \
--target_shape_name=conv_without_bias \
--max_steps=500000 \
--hidden_dim=128 \
--num_cold_steps=10000 \
--save_checkpoint_for_resuming_training \
--use_wandb

推理

将输入图像放入文件夹中。
预处理(如果需要)。对于人脸领域,如果输入图像尚未裁剪和对齐,请运行以下脚本预处理数据。

RAW_IMAGE_DIR=""
PROCESSED_IMAGE_DIR=""

python scripts/align_all_parallel.py \
--raw_dir "$RAW_IMAGE_DIR" \
--saved_dir "$PROCESSED_IMAGE_DIR" \
--num_threads 8

参数说明如下。

参数	描述
`RAW_IMAGE_DIR`	包含原始输入图像的文件夹路径
`PROCESSED_IMAGE_DIR`	保存处理后输入图像的文件夹路径

运行推理

在运行以下脚本之前正确设置参数。

INPUT_DATA_DIR=""
RESULT_DIR=""
MODEL_PATH=""
GPU_ID=0

CUDA_VISIBLE_DEVICES="$GPU_ID" \
python scripts/inference.py \
--exp_dir="$RESULT_DIR" \
--checkpoint_path="$MODEL_PATH" \
--data_path="$INPUT_DATA_DIR" \
--batch_size=4 \
--workers=4

参数说明如下。

参数	描述
`RESULT_DIR`	保存推理结果的目录
`MODEL_PATH`	HyperInverter模型的路径
`INPUT_DATA_DIR`	包含处理后输入图像的文件夹路径

最后,重建的图像可以在RESULT_DIR/inference_results文件夹中找到。

定量评估

我们已经准备了人脸和教堂模型的定量重建评估脚本。请在这些文件中设置与您的模型兼容的参数,并运行以下命令进行评估过程。

人脸

sh sample_scripts/human_faces_reconstruction_quantitative_evaluation.sh

教堂

sh sample_scripts/church_reconstruction_quantitative_evaluation.sh

定性比较

对于以下实验,请访问文件configs/paths_config.py并使用HyperInverter和其他反演方法的预训练模型的路径更新model_paths字典。对于其他反演方法,请访问它们的Github仓库下载预训练权重。

重建

定性重建比较的示例脚本是:

DOMAIN=""
METHODS=""
INPUT_DATA_DIR=""
SAVED_RESULTS_DIR_NAME=""
MAX_NUM_IMAGES=100
SAVED_SIZE=1024
GPU_ID=0


CUDA_VISIBLE_DEVICES="$GPU_ID" \
python evaluation/reconstruction_comparison.py \
--methods="$METHODS" \
--domain="$DOMAIN" \
--input_data_dir="$INPUT_DATA_DIR"  \
--input_data_id="$SAVED_RESULTS_DIR_NAME" \
--output_dir=outputs \
--saved_embedding_dir=embeddings \
--max_num_images="$MAX_NUM_IMAGES" \
--resize="$SAVED_SIZE"

参数说明如下。

参数	描述
`DOMAIN`	输入域,选项为{`human_faces`, `churches`}
`METHODS`	反演方法,用逗号分隔,支持的方法有{`hyper_inverter`, `psp`, `e4e`, `SG2_plus`, `SG2`, `w_encoder`, `hyper_inverter`, `restyle_e4e`}
`INPUT_DATA_DIR`	包含处理后输入图像的文件夹路径
`SAVED_RESULTS_DIR_NAME`	保存结果的文件夹名称
`MAX_NUM_IMAGES`	要处理的最大图像数量
`SAVED_SIZE`	每种方法保存的图像大小

结果可以在outputs/SAVED_RESULTS_DIR_NAME文件夹中找到。

编辑

编辑比较的示例脚本是:

DOMAIN=""
METHODS=""
DIRECTION=""
INPUT_DATA_DIR=""
SAVED_RESULTS_DIR_NAME=""
MAX_NUM_IMAGES=10
SAVED_SIZE=1024
MIN_MAG=-30  
MAX_MAG=30
STEP=5 
GPU_ID=0

CUDA_VISIBLE_DEVICES="$GPU_ID" \
python evaluation/editing_inference.py \
--methods="$METHODS" \
--domain="$DOMAIN" \
--input_data_dir="$INPUT_DATA_DIR" \
--input_data_id="$SAVED_RESULTS_DIR_NAME" \
--output_dir=outputs \
--saved_embedding_dir=embeddings \
--direction="$DIRECTION" \
--min_factor="$MIN_MAG" \
--max_factor="$MAX_MAG" \
--step="$STEP" \
--max_num_images="$MAX_NUM_IMAGES" \
--resize="$SAVED_SIZE" \
--save_edited_images \
--gif_speed=4

结果可以在outputs/SAVED_RESULTS_DIR_NAME文件夹中找到。请尝试不同的MIN_MAG和MAX_MAG值以获得最佳结果。参数说明如下所示。

参数	描述
`DOMAIN`	输入域，选项包括{`human_faces`, `churches`}
`METHODS`	反演方法，用逗号分隔，支持的方法有{`hyper_inverter`, `psp`, `e4e`, `SG2_plus`, `SG2`, `w_encoder`, `restyle_e4e`}，例如：`hyper_inverter,psp,e4e`
`DIRECTION`	编辑方向，支持的方向见下表
`INPUT_DATA_DIR`	包含已处理输入图像的文件夹路径
`MIN_MAG`	最小编辑幅度，请调整此参数以获得最佳结果
`MAX_MAG`	最大编辑幅度，请调整此参数以获得最佳结果
`STEP`	从最小幅度到最大幅度的步长
`SAVED_RESULTS_DIR_NAME`	保存结果的文件夹名称
`MAX_NUM_IMAGES`	要处理的最大图像数量
`SAVED_SIZE`	每种方法保存图像的大小

支持的编辑方向如下所列。

人脸

方法	编辑方向
GANSpace	`eye_openness`, `trimmed_beard`, `lipstick` `face_roundness`, `nose_length`, `eyebrow_thickness`, `head_angle_up`, `displeased`
InterFaceGAN	`age`, `smile`, `rotation`
StyleCLIP	`surprised`, `afro`, `angry`, `beyonce`, `bobcut`, `bowlcut`, `curly_hair`, `hilary_clinton`, `depp`, `mohawk`, `purple_hair`, `taylor_swift`, `trump`, `zuckerberg`

教堂

方法	编辑方向
GANSpace	`clouds`, `vibrant`, `blue_skies`, `trees`

应用：真实图像插值

两张真实图像插值的示例脚本如下：

DOMAIN=""
METHODS=""
PATH_TO_INPUT_IMAGE_1=""
PATH_TO_INPUT_IMAGE_2=""
SAVED_RESULTS_DIR=""
SAVED_RESULTS_FILE_NAME=""
SAVED_SIZE=1024
NUM_STEPS=100
GPU_ID=0

CUDA_VISIBLE_DEVICES="$GPU_ID" \
python evaluation/real_image_interpolation.py \
--domain="$DOMAIN" \
--method="$METHODS" \
--left_image_path="$PATH_TO_INPUT_IMAGE_1" \
--right_image_path="$PATH_TO_INPUT_IMAGE_2" \
--steps="$NUM_STEPS" \
--saved_image_size="$SAVED_SIZE" \
--saved_dir="$SAVED_RESULTS_DIR" \
--saved_file_name="$SAVED_RESULTS_FILE_NAME" \
--save_interpolated_images \
--gif_speed=2

参数说明如下所示。

参数	描述
`DOMAIN`	输入域，选项包括{`human_faces`, `churches`}
`METHODS`	反演方法，用逗号分隔，支持的方法有{`hyper_inverter`, `psp`, `e4e`, `SG2_plus`, `SG2`, `w_encoder`, `restyle_e4e`}，例如：`hyper_inverter,psp,e4e`
`PATH_TO_INPUT_IMAGE_1`	输入图像1的路径
`PATH_TO_INPUT_IMAGE_2`	输入图像2的路径
`NUM_STEPS`	插值步数
`SAVED_RESULTS_DIR`	保存结果的文件夹路径
`MAX_NUM_IMAGES`	要处理的最大图像数量
`SAVED_RESULTS_FILE_NAME`	GIF结果文件的名称
`SAVED_SIZE`	每种方法保存图像的大小