自注意力生成对抗网络

Han Zhang、Ian Goodfellow、Dimitris Metaxas 和 Augustus Odena，《自注意力生成对抗网络》。arXiv 预印本 arXiv:1805.08318 (2018)。

项目概述

本仓库提供了 SAGAN 的 PyTorch 实现。WGAN-GP 和 WGAN-Hinge 损失均已实现，但请注意，WGAN-GP 与谱归一化并不兼容。若要使用 WGAN-GP，则需移除模型中的所有谱归一化层。

自注意力机制被应用到判别器和生成器的最后两层。

当前更新状态

• 有监督设置
• TensorBoard 日志记录
• [20180608] 更新了自注意力模块。感谢我的同事 Cheonbok Park！请参阅 'sagan_models.py' 中的更新内容。该实现应更加高效，并支持大尺寸图像的训练
• 注意力可视化（LSUN 教堂-户外场景）
• 无监督设置（暂未使用标签）
• 应用：谱归一化，代码来自 这里
• 实现：自注意力模块、双时间尺度更新规则（TTUR）、WGAN-Hinge 损失、WGAN-GP 损失

   

结果

LSUN 数据集上的注意力结果（第8轮迭代）

这是 SAGAN 在 LSUN 教堂-户外场景数据集上的像素级注意力结果。尽管仅凭注意力图本身难以解释，但无监督训练下的自注意力模块仍然有效。未来将补充更多关于生成图像的效果展示。这些是生成器第3层和第4层的自注意力可视化，分别对应16×16和32×32大小的特征图，每张图包含64张图像。为了便于可视化，仅选择了部分像素点进行标注，如最左侧和最右侧所示。

CelebA 数据集（左侧为当前轮次，仍在训练中）

LSUN 教堂-户外场景数据集（左侧为当前轮次，仍在训练中）

前置条件

• Python 3.5+
• PyTorch 0.3.0

 

使用方法

1. 克隆仓库

$ git clone https://github.com/heykeetae/Self-Attention-GAN.git
$ cd Self-Attention-GAN

2. 安装数据集（CelebA 或 LSUN）

$ bash download.sh CelebA
或
$ bash download.sh LSUN

3. 训练

(i) 训练

$ python main.py --batch_size 64 --imsize 64 --dataset celeb --adv_loss hinge --version sagan_celeb
或
$ python main.py --batch_size 64 --imsize 64 --dataset lsun --adv_loss hinge --version sagan_lsun

4. 查看结果

$ cd samples/sagan_celeb
或
$ cd samples/sagan_lsun

每100次迭代会生成一批样本。采样频率可通过 --sample_step 参数控制（例如，--sample_step 100）。

Self-Attention-GAN 快速上手指南

本指南基于 Han Zhang 等人提出的自注意力生成对抗网络（SAGAN）PyTorch 实现，帮助开发者快速在本地运行模型并生成图像。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux / macOS / Windows
• Python: 3.5 或更高版本
• 深度学习框架: PyTorch 0.3.0 (注意：原文指定版本较老，若使用新版 PyTorch 可能需要微调代码)
• 依赖库: 建议安装 numpy, scipy, pillow 等常用科学计算库

提示：国内用户可使用清华源或阿里源加速 Python 包的安装：

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

(注：若仓库根目录无 requirements.txt，请手动安装上述依赖)

安装步骤

1. 克隆项目代码

将仓库克隆到本地并进入项目目录：

git clone https://github.com/heykeetae/Self-Attention-GAN.git
cd Self-Attention-GAN

2. 下载数据集

项目支持 CelebA 和 LSUN 数据集。使用提供的脚本自动下载（请确保网络连接正常，国内用户若下载缓慢可考虑手动下载后放入对应目录）：

下载 CelebA 数据集：

bash download.sh CelebA

或者下载 LSUN 数据集：

bash download.sh LSUN

基本使用

完成环境配置和数据下载后，即可开始训练模型。以下是最简单的训练命令示例。

训练模型

在 CelebA 数据集上训练：

python main.py --batch_size 64 --imsize 64 --dataset celeb --adv_loss hinge --version sagan_celeb

在 LSUN 数据集上训练：

python main.py --batch_size 64 --imsize 64 --dataset lsun --adv_loss hinge --version sagan_lsun

参数说明：

• --batch_size: 批处理大小（默认示例为 64）
• --imsize: 图像分辨率（默认示例为 64）
• --adv_loss: 对抗损失类型，支持 hinge (WGAN-hinge) 或 gp (WGAN-GP)。注意：若使用 gp 模式，需移除模型中的谱归一化（Spectral Normalization）。
• --version: 实验版本标识，用于区分保存路径。

查看生成结果

训练过程中，模型会每隔一定迭代次数自动生成样本图片。

• CelebA 结果路径: samples/sagan_celeb
• LSUN 结果路径: samples/sagan_lsun

您可以通过 --sample_step 参数控制采样频率（例如每 100 次迭代保存一次）：

python main.py --batch_size 64 --imsize 64 --dataset celeb --adv_loss hinge --version sagan_celeb --sample_step 100

训练完成后，直接进入上述样本目录即可查看生成的图像效果。