GANformer：生成对抗变换器

德鲁·A·哈德森 & C·劳伦斯·齐特尼克

欢迎查看我们的全新PyTorch版本以及GANformer2论文！

更新（2022年2月21日）： 我们已将PyTorch版本的权重初始化调整为预期规模，从而显著提升了模型的学习速度！

这是GANformer模型的实现，该模型是一种新颖且高效的变换器架构，专为图像生成任务而设计。网络采用一种_二分结构_，能够在保持线性计算效率的同时，实现跨图像的长距离交互，并轻松扩展到高分辨率合成。

该模型通过迭代地在潜在变量集合与不断演化的视觉特征之间传递信息，使两者相互补充、相互促进，从而鼓励对象和场景的组合式表征的出现。 与经典的变换器架构不同，它使用乘法融合机制，允许灵活的区域调制，因此可以被视为对成功StyleGAN网络的一种推广。

第一篇论文：https://arxiv.org/pdf/2103.01209
第二篇论文：https://arxiv.org/abs/2111.08960
联系方式：dorarad@cs.stanford.edu
实现代码：network.py (TF / Pytorch)

我们现在同时支持PyTorch和TF！

:white_check_mark: 上传初始代码和README文件
:white_check_mark: 图像采样与可视化脚本
:white_check_mark: 代码清理与重构，添加文档
:white_check_mark: 训练与数据预处理说明
:white_check_mark: 所有数据集的预训练网络
:white_check_mark: 额外的可视化与评估
:white_check_mark: 提供训练时间更长的模型
:white_check_mark: 发布PyTorch版本
:white_check_mark: 发布高分辨率（最高1024×1024）的预训练模型
⬜️ 发布GANformer2模型（支持布局生成及条件布局到图像生成）

如果您遇到任何问题，或有任何改进建议及功能扩展需求，请随时通过Issues页面或发送邮件至dorarad@stanford.edu与我联系。

Bibtex

@article{hudson2021ganformer,
  title={Generative Adversarial Transformers},
  author={Hudson, Drew A and Zitnick, C. Lawrence},
  journal={Proceedings of the 38th International Conference on Machine Learning, {ICML} 2021},
  year={2021}
}

@article{hudson2021ganformer2,
  title={Compositional Transformers for Scene Generation},
  author={Hudson, Drew A and Zitnick, C. Lawrence},
  journal={Advances in Neural Information Processing Systems {NeurIPS} 2021},
  year={2021}
}

样例图片

使用预训练模型（这些模型仅需比StyleGAN2少5-7倍的训练步数即可生成！进一步延长训练时间将进一步提升图像质量）：

环境要求

• 支持Python 3.6或3.7。
• 对于TF版本：我们推荐使用开发时所用的TensorFlow 1.14，但TensorFlow 1.15也同样适用。
• 对于PyTorch版本：我们支持PyTorch >= 1.8。
• 代码已在CUDA 10.0工具包和cuDNN 7.5环境下测试通过。
• 我们曾在Titan V GPU上进行过实验。假设拥有12GB显存（更多显存可加速训练）。
• 请参阅requirements.txt（TF / PyTorch），其中列出了所需的Python包，运行pip install -r requirements.txt即可安装。

快速入门与概览

我们的仓库同时支持TensorFlow（主目录下）和PyTorch（位于pytorch_version）。两种实现的代码和文件结构非常接近，接口也完全一致。如需从TF切换至PyTorch，只需进入pytorch_version，并安装其中的依赖项。 如有任何关于新实现的问题或建议，欢迎随时提交Issue或直接联系[dorarad@cs.stanford.edu]！

一个使用预训练GANformer的最小示例可在generate.py中找到（TF / PyTorch）。执行该程序时，这段仅有10行的代码会下载一个预训练模型，并用其生成一些图像：

python generate.py --gpus 0 --model gdrive:bedrooms-snapshot.pkl --output-dir images --images-num 32

您可以通过--truncation-psi来控制生成图像的质量与多样性之间的权衡。
我们建议尝试0.6-1.0范围内的不同值。

预训练模型与高分辨率

我们为所有数据集提供了256×256分辨率的预训练模型，同时为FFHQ提供了1024×1024分辨率的模型，为Cityscapes提供了1024×2048分辨率的模型。

要生成高分辨率模型的图像，请运行以下命令： （我们将批量大小降低到1，以便它们可以加载到单个GPU上）

python generate.py --gpus 0 --model gdrive:ffhq-snapshot-1024.pkl --output-dir ffhq_images --images-num 32 --batch-size 1
python generate.py --gpus 0 --model gdrive:cityscapes-snapshot-2048.pkl --output-dir cityscapes_images --images-num 32 --batch-size 1 --ratio 0.5 # 目前仅TF版本支持1024 x 2048的Cityscapes

我们可以使用run_network.py（TF / Pytorch）对新模型或预训练模型进行定量和定性的训练与评估。
模型架构可以在network.py（TF / Pytorch）中找到。训练流程则在training_loop.py（TF / Pytorch）中实现。

数据准备

我们在四个图像和场景数据集上探索了GANformer模型：CLEVR、LSUN-Bedrooms、Cityscapes和FFHQ。该模型也可以在其他数据集上进行训练。 我们以256×256分辨率训练了模型。更高分辨率同样受支持。模型会自动适应数据集中图像的分辨率。

prepare_data.py（TF / Pytorch）既可以从我们的目录中准备数据集，也可以创建新的数据集。

默认数据集

要从目录中准备数据集，请运行以下命令：

python prepare_data.py --ffhq --cityscapes --clevr --bedrooms --max-images 100000

下表详细列出了目录中的数据集信息。

常用选项：

• --data-dir 输出数据目录（默认：datasets）
• --shards-num 选择数据分片的数量（默认：根据每个数据集调整）
• --max-images 只存储数据集的一部分，以减少存储的tfrecord/图像文件大小（默认：最大）。
  这对于节省大型数据集的空间特别有用，例如LSUN-bedrooms（原始包含300万张图像）。

自定义数据集

您也可以使用该脚本创建新的自定义数据集。例如：

python prepare_data.py --task <dataset-name> --images-dir <source-dir> --format png --ratio 0.7 --shards-num 5

该脚本支持多种格式：png、jpg、npy、hdf5、tfds和lmdb。

数据集目录

数据集	图像数量	分辨率	下载大小	TFrecords大小	Gamma
FFHQ	70,000	256×256	13GB	13GB	10
CLEVR	100,015	256×256	18GB	15.5GB	40
Cityscapes	24,998	256×256	1.8GB	8GB	20
LSUN-Bedrooms	3,033,042	256×256	42.8GB	最高可达480GB	100

使用--max-images可以减小tfrecord文件的大小。

训练

模型通过使用--train选项进行训练。要微调一个预训练的GANformer模型：

python run_network.py --train --gpus 0 --ganformer-default --expname clevr-pretrained --dataset clevr \
  --pretrained-pkl gdrive:clevr-snapshot.pkl

我们为bedrooms、cityscapes、clevr和ffhq提供了预训练模型。

要从头开始以默认配置训练GANformer模型：

python run_network.py --train --gpus 0 --ganformer-default --expname clevr-scratch --dataset clevr --eval-images-num 10000

默认情况下，模型训练会从最新的快照恢复。使用--restart可开始一个新的实验，或使用--pretrained-pkl选择特定的快照进行加载。

为了与当前最先进的方法进行比较，我们使用50,000个样本图像计算指标分数。然而，为了加快训练速度，我们建议将--eval-images-num设置为较低的数值。请注意，这可能会影响指标的精确性，因此我们建议在训练过程中使用较低的值，并在最终评估时将其提高。

我们支持多种命令行选项来调整模型、训练和评估。运行python run_network.py -h即可查看完整的选项列表！

我们建议在训练新数据集时尝试不同的--gamma值。如果您以≥512分辨率进行训练并遇到OOM问题，请考虑将--batch-gpu降低到较低的值。

日志记录

• 在训练过程中，样本图像和注意力图会被生成并存储在results/<expname>-<run-id>目录下（--keep-samples）。
• 指标也会定期计算并记录在metric-<name>.txt文件中。--metrics可设置为fid用于FID、is用于Inception Score以及pr用于Precision/Recall。
• 还会创建TensorBoard日志（--summarize），跟踪指标、生成器和判别器的损失值以及其他有用的统计信息。

基线模型

除了GANformer之外，代码库还支持多种基线模型。例如，要运行一个普通的GAN模型：

python run_network.py --train --gpus 0 --baseline GAN --expname clevr-gan --dataset clevr 

• Vanilla GAN：--baseline GAN，一种没有风格调制的标准GAN。
• StyleGAN2：--baseline StyleGAN2，使用一个全局潜在变量来调制图像特征。
• k-GAN：--baseline kGAN，独立生成多个图像层，然后将它们合并为一张共享图像（仅TF版本支持）。
• SAGAN：--baseline SAGAN，在低分辨率层（如32x32）中对所有图像特征进行自注意力操作（仅TF版本支持）。

评估

要评估模型，请使用 --eval 选项：

python run_network.py --eval --gpus 0 --expname clevr-exp --dataset clevr

添加 --pretrained-pkl gdrive:<dataset>-snapshot.pkl 可以评估预训练模型。

下面我们提供了 GANformer（使用上述预训练检查点）以及基线模型的 FID-50k 分数。
请注意，这些分数与 StyleGAN2 论文中报告的分数不同，因为我们的实验进行了多达 7 倍的训练步数（所有模型的实验中为 5k–15k kimg 步，使用 4 张 GPU 卡大约需要 3–4 天），而他们的实验则进行了 50k–70k kimg 步，耗时超过 90 个 GPU 天。

关于生成器/判别器的说明：经过消融实验，我们发现将单纯形和双工注意力机制仅应用于生成器（而非同时应用于生成器和判别器）能够提升模型性能。因此，我们发布了仅在生成器中加入注意力机制的预训练模型，并已更新论文以反映这一改进！

模型	CLEVR	LSUN-Bedroom	FFHQ	Cityscapes
GAN	25.02	12.16	13.18	11.57
kGAN	28.28	69.9	61.14	51.08
SAGAN	26.04	14.06	16.21	12.81
StyleGAN2	16.05	11.53	16.21	8.35
GANformer	9.24	6.15	7.42	5.23

模型变更日志

与论文中描述的原始 GANformer 相比，本仓库进行了多项额外改进，从而提升了性能：

• 使用 --mapping_ltnt2ltnt 选项，使潜在变量在映射网络内部通过自注意力机制直接相互交流，然后再开始生成图像。
• 在 k 个潜在组件之外增加一个全局潜在变量 (--style)。首先由全局潜在变量统一调节图像的所有特征，然后由 k 个潜在变量根据二分图变换器的注意力机制分别调节不同区域。
  全局潜在变量有助于协调图像的整体方面，例如全局光照条件、人脸等的整体风格属性等。
• 经过这些调整后，我们发现向判别器添加变换器层并无额外益处，因此为了简化，我们将其禁用了。

可视化

该代码支持生成定性结果和可视化效果。例如，要创建每一层的注意力图：

python run_network.py --gpus 0 --vis --expname clevr-exp --dataset clevr --vis-layer-maps

下面可以看到 GANformer 生成的示例图像和注意力图：

命令行选项

以下列出了一些最常用的模型选项。

训练

• --gamma：建议针对所选数据集尝试不同的值（默认值为 10）。
• --truncation-psi：控制图像质量与多样性之间的权衡。（默认值为 0.7）
• --eval-images-num：用于计算指标的图像数量。建议选择较小的数值以加快训练速度（默认值为 50,000）。
• --restart：从头开始重新训练，而不是从最新快照恢复。
• --pretrained-pkl：加载预训练模型，可以是本地文件，也可以是从云端加载 gdrive:<dataset>-snapshot.pkl 格式的文件，适用于目录中的各类数据集。
• --data-dir 和 --result-dir：分别指定数据集（tfrecords）和日志/结果的存储目录。

模型（最常用）

• --transformer：在生成器中添加变换器层（GANformer）。
• --components-num：潜在组件的数量，这些组件将对图像进行注意力操作。建议取值范围为 8–16（默认值为 1）。
• --latent-size：整体潜在空间大小（默认值为 512）。每个潜在组件的大小则为 latent_size/components_num。
• --num-heads：注意力头的数量（默认值为 1）。
• --integration：变换器层中信息的整合方式，例如 add 或 mul（默认值为 mul）。

模型（其他）

• --g-start-res 和 --g-end-res：变换器层的起始和结束分辨率（默认值为所有层直至分辨率为 2⁸）。
• --kmeans：跟踪并更新图像到潜在变量的分配中心点，用于双工注意力机制。
• --mapping-ltnt2ltnt：在映射网络中对潜在变量执行自注意力操作。
• --use-pos：为潜在变量使用可训练的位置编码。
• --style False：关闭单向量全局风格调节功能（StyleGAN2）。

可视化

• 样本图像
  • --vis-images：生成图像样本。
  • --vis-latents：保存源潜在向量。
• 注意力图
  • --vis-maps：可视化最后一层和第一个注意力头的注意力图。
  • --vis-layer-maps：可视化所有层和所有注意力头的注意力图。
  • --blending-alpha：在可视化图像与注意力图的混合时设置透明度权重。
• 图像插值
  • --vis-interpolations：生成源潜在变量之间的一系列插值图像。
  • --interpolation-density：插值过程中两个端点之间的样本数量（默认值为 8）。
• 其他
  • --vis-noise-var：创建噪声变化的可视化效果。
  • --vis-style-mix：创建风格混合的可视化效果。

运行 python run_network.py -h 可查看完整的选项列表。

示例图像（更多案例）

CUDA / 安装

该模型依赖于自定义的 TensorFlow/PyTorch 操作，这些操作使用 NVCC 动态编译。

要设置环境，例如针对 CUDA 10.0：

export PATH=/usr/local/cuda-10.0/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda10.0/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

要测试 NVCC 安装是否正常工作，请运行：

nvcc test_nvcc.cu -o test_nvcc -run
| CPU 说你好。
| GPU 说你好。

在 PyTorch 版本中，如果出现以下重复信息： “无法构建 upfirdn2d 的 CUDA 内核。正在回退到低效的参考实现。” 请确保你的 CUDA 和 PyTorch 版本匹配。如果你安装了多个 CUDA 版本，可以考虑将 CUDA_HOME 设置为对应的版本。例如：

export CUDA_HOME=/usr/local/cuda-10.1

架构概述

GANformer 由两个网络组成：

生成器：根据随机采样的潜在变量 (z) 生成图像 (x)。潜在变量 z 的形状为 [batch_size, component_num, latent_dim]，其中 component_num 默认为 1（Vanilla GAN、StyleGAN），但在 GANformer 模型中大于 1。我们可以通过沿第二个维度拆分 z 来获得 z_1,...,z_k 个潜在组件。生成器同样分为两部分：

• 映射网络：将来自正态分布的采样潜在变量 (z) 转换为中间空间 (w)。它由一系列前馈层组成。这 k 个潜在组件可以从 z 空间独立映射到 w 空间，也可以通过自注意力机制相互作用（可选标志）。
• 合成网络：中间潜在变量 w 用于指导新图像的生成。图像特征从一个小型的常量或采样网格（4x4）开始，然后经过多层卷积和上采样，直到达到所需的分辨率（如 256x256）。每次卷积后，图像特征都会被中间潜在向量 w 调制（即控制其方差和偏置）。而在 StyleGAN 模型中，只有一个全局的 w 向量来平等地控制所有特征。GANformer 使用注意力机制，使 k 个潜在组件分别专注于控制图像中的不同区域，从而协同生成图像，因此在生成包含多物体场景的图像时表现更佳。
• 注意力机制可以以多种方式应用：
  • 单向注意力：当注意力仅从潜在变量单向作用于图像特征时（自上而下）。
  • 双向注意力：当注意力同时在两个方向上作用——从潜在变量到图像特征（自上而下），再从图像特征反馈回潜在变量（自下而上），从而使两种表示相互迭代地提供信息。
  • 潜在变量之间的自注意力：也可用于直接促进潜在变量之间的交互。
  • 图像特征之间的自注意力（SAGAN 模型）：早期方法直接在图像特征之间使用注意力，但这种方法扩展性较差，因为特征数量与分辨率成平方关系，在高分辨率下会变得非常高。

判别器：接收一张图像，并预测它是真实的还是伪造的——即来源于数据集还是生成器。该模型会对图像进行多层卷积和下采样，逐步降低图像表示的分辨率，直至做出最终预测。可选地，判别器中也可以加入注意力机制，配备多个（k）聚合变量，这些变量利用注意力机制在图像处理过程中自适应地收集信息。我们观察到，在判别器中使用注意力机制时，模型性能会有小幅提升，不过需要注意的是，根据我们的观察，使用注意力机制所带来的主要收益来自于生成器。

代码库

该代码库基于并扩展了 Karras 等人优秀的 StyleGAN2 和 StyleGAN2-ADA 仓库。

GANformer 模型也可以被视为 StyleGAN 的一种泛化：StyleGAN 只有一个全局潜在向量来全局控制所有图像特征的风格，而 GANformer 则拥有 k 个潜在向量，它们通过注意力机制协作控制图像中的不同区域，从而更好地建模多物体和组合场景的图像。

致谢

我谨向 Christopher D. Manning 表示感谢，感谢他在开发二部图 Transformer 时给予的富有成效的讨论和建设性意见，尤其是在语言表示领域内的探索；同时也感谢他提供的慷慨资金支持，使得这项工作得以顺利开展！ :sunflower:

如果您有任何问题、评论或反馈，请随时通过 dorarad@cs.stanford.edu 与我联系。谢谢！ :)

GANformer 快速上手指南

GANformer 是一种用于图像生成的新型高效 Transformer 模型。它采用二分结构，支持长距离交互，同时保持线性计算效率，可扩展至高分辨率合成。本项目同时支持 TensorFlow 和 PyTorch 版本。

1. 环境准备

系统要求

• 操作系统: Linux (推荐)
• GPU: 建议显存 ≥ 12GB (如 Titan V)，支持 CUDA 10.0+
• Python: 3.6 或 3.7

框架依赖

根据你选择的版本安装对应框架：

• TensorFlow 版本: TensorFlow 1.14 (推荐) 或 1.15
• PyTorch 版本: PyTorch ≥ 1.8
• 其他: cuDNN 7.5+

国内加速提示: 安装 Python 包时建议使用清华或阿里镜像源以加快下载速度。

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装步骤

本项目主目录为 TensorFlow 版本，pytorch_version 目录为 PyTorch 版本。以下以主目录（TF）为例，PyTorch 用户请进入 pytorch_version 目录操作。

1. 克隆代码库

   git clone <repository_url>
   cd gansformer
   # 如果使用 PyTorch 版本，请执行: cd pytorch_version
   

2. 安装依赖包

   pip install -r requirements.txt
   

   (注：PyTorch 用户请运行 pip install -r pytorch_version/requirements.txt)

3. 验证安装 确保 python 命令可用，且 GPU 驱动正常识别。

3. 基本使用

快速生成图像 (使用预训练模型)

最简单的用法是运行 generate.py 脚本。该脚本会自动下载预训练模型并生成图像。

命令示例：

python generate.py --gpus 0 --model gdrive:bedrooms-snapshot.pkl --output-dir images --images-num 32

参数说明：

• --gpus: 指定使用的 GPU ID (例如 0)。
• --model: 预训练模型路径。支持从 Google Drive 自动下载 (格式 gdrive:<name>.pkl) 或本地路径。
  • 可用模型包括：bedrooms-snapshot.pkl, ffhq-snapshot.pkl, clevr-snapshot.pkl, cityscapes-snapshot.pkl。
• --output-dir: 生成图像的保存目录。
• --images-num: 生成图像的数量。

调整生成质量与多样性： 使用 --truncation-psi 参数控制权衡，推荐范围 0.6 到 1.0。

python generate.py --gpus 0 --model gdrive:ffhq-snapshot.pkl --output-dir ffhq_images --images-num 32 --truncation-psi 0.7

生成高分辨率图像

对于 1024x1024 (FFHQ) 或 1024x2048 (Cityscapes) 的模型，需将批大小设为 1 以适应显存：

# FFHQ 1024x1024
python generate.py --gpus 0 --model gdrive:ffhq-snapshot-1024.pkl --output-dir ffhq_images --images-num 32 --batch-size 1

# Cityscapes 1024x2048 (目前仅 TF 版本支持)
python generate.py --gpus 0 --model gdrive:cityscapes-snapshot-2048.pkl --output-dir cityscapes_images --images-num 32 --batch-size 1 --ratio 0.5

数据准备 (可选)

如需从头训练，可使用 prepare_data.py 准备数据集（支持 CLEVR, LSUN-Bedrooms, Cityscapes, FFHQ）。

下载并准备默认数据集：

python prepare_data.py --ffhq --cityscapes --clevr --bedrooms --max-images 100000

提示：使用 --max-images 可限制下载数量以节省空间。

开始训练 (可选)

使用 run_network.py 进行训练或微调。

微调预训练模型：

python run_network.py --train --gpus 0 --ganformer-default --expname clevr-finetune --dataset clevr --pretrained-pkl gdrive:clevr-snapshot.pkl

从头训练：

python run_network.py --train --gpus 0 --ganformer-default --expname clevr-scratch --dataset clevr --eval-images-num 10000

训练日志、样本图像及 Tensorboard 记录将保存在 results/<expname>-<run-id> 目录下。