PyTorch 预训练 BigGAN

基于 DeepMind 的 BigGAN 模型，并附带 DeepMind 提供的预训练权重，以操作对操作的方式在 PyTorch 中重新实现。

简介

本仓库包含 DeepMind 在论文《用于高保真自然图像合成的大规模 GAN 训练》（作者：Andrew Brock、Jeff Donahue 和 Karen Simonyan）中发布的 BigGAN 的 PyTorch 版本，该版本实现了与 TensorFlow 原版相同的操作逻辑。

此 PyTorch 实现附带了 DeepMind 提供的 128x128、256x256 和 512x512 分辨率的预训练模型。我们还提供了用于从 TensorFlow Hub 下载并转换这些模型的脚本。

该重实现基于 TensorFlow 版本的原始计算图完成，其行为与 TensorFlow 版本高度一致，输出差异的数量级约为 1e-5。

目前，该实现仅包含生成器部分，因为判别器的权重并未公开（尽管判别器的结构与生成器非常相似，因此可以很容易地添加）。如果您希望为此提交 PR，我将非常乐意提供帮助。

安装

本仓库已在 Python 3.6 和 PyTorch 1.0.1 上测试通过。

可以通过 pip 安装 PyTorch 预训练 BigGAN：

pip install pytorch-pretrained-biggan

如果您只是想简单体验一下这个 GAN，那么以上安装步骤就足够了。

如果您还想使用转换脚本和 ImageNet 工具，则需要额外的依赖，尤其是 TensorFlow 和 NLTK。要安装所有依赖，请使用 full_requirements.txt 文件：

git clone https://github.com/huggingface/pytorch-pretrained-BigGAN.git
cd pytorch-pretrained-BigGAN
pip install -r full_requirements.txt

模型

本仓库提供了对 BigGAN “deep” 版本的直接且简便访问，支持 128、256 和 512 像素分辨率，具体如相关论文所述。以下是各模型的详细信息：

• BigGAN-deep-128：一个拥有 5040 万参数、生成 128x128 像素图像的模型，模型文件大小为 201 MB；
• BigGAN-deep-256：一个拥有 5590 万参数、生成 256x256 像素图像的模型，模型文件大小为 224 MB；
• BigGAN-deep-512：一个拥有 5620 万参数、生成 512x512 像素图像的模型，模型文件大小为 225 MB。

有关架构的详细信息，请参阅论文的附录 B。

所有模型都预先计算好了 51 种截断值（范围从 0 到 1）对应的批量归一化统计信息，详情请参阅论文附录 C.1。

使用

以下是一个使用预训练模型的 BigGAN 快速入门示例。

有关这些类和方法的详细信息，请参阅下方的 文档部分。

import torch
from pytorch_pretrained_biggan import (BigGAN, one_hot_from_names, truncated_noise_sample,
                                       save_as_images, display_in_terminal)

# 可选：如果想了解更多信息，可以启用日志记录
import logging
logging.basicConfig(level=logging.INFO)

# 加载预训练模型分词器（词汇表）
model = BigGAN.from_pretrained('biggan-deep-256')

# 准备输入
truncation = 0.4
class_vector = one_hot_from_names(['肥皂泡', '咖啡', '蘑菇'], batch_size=3)
noise_vector = truncated_noise_sample(truncation=truncation, batch_size=3)

# 将所有数据转换为张量
noise_vector = torch.from_numpy(noise_vector)
class_vector = torch.from_numpy(class_vector)

# 如果有 GPU，将所有数据转移到 CUDA 上
noise_vector = noise_vector.to('cuda')
class_vector = class_vector.to('cuda')
model.to('cuda')

# 生成图像
with torch.no_grad():
    output = model(noise_vector, class_vector, truncation)

# 如果有 GPU，再将结果移回 CPU
output = output.to('cpu')

# 如果终端支持 sixel 格式，可以直接在终端中显示图像
# （详情请参见 https://github.com/saitoha/libsixel）
display_in_terminal(output)

# 将结果保存为 PNG 图像
save_as_images(output)

文档

加载 DeepMind 的预训练权重

要加载 DeepMind 的预训练模型之一，可以使用 from_pretrained() 方法实例化一个 BigGAN 模型，如下所示：

model = BigGAN.from_pretrained(PRE_TRAINED_MODEL_NAME_OR_PATH, cache_dir=None)

其中：

• PRE_TRAINED_MODEL_NAME_OR_PATH 可以是以下两种形式之一：
  • Google AI 或 OpenAI 提供的预训练模型的快捷名称，可从以下列表中选择：
    • biggan-deep-128：12 层，768 隐藏单元，12 头注意力机制，1.1 亿参数；
    • biggan-deep-256：24 层，1024 隐藏单元，16 头注意力机制，3.4 亿参数；
    • biggan-deep-512：12 层，768 隐藏单元，12 头注意力机制，1.1 亿参数。
  • 预训练模型存档的路径或 URL，其中包含：
    • config.json：模型配置文件；
    • pytorch_model.bin：BigGAN 预训练实例的 PyTorch 模型文件（使用常规的 torch.save() 保存）。 如果 PRE_TRAINED_MODEL_NAME_OR_PATH 是快捷名称，预训练权重将从 AWS S3 下载（详见 此处 的链接），并存储在缓存目录中以避免重复下载（默认缓存目录为 ~/.pytorch_pretrained_biggan/）。
• cache_dir 是一个可选参数，用于指定下载和缓存预训练模型权重的具体目录。

配置

BigGANConfig 是一个用于存储和加载 BigGAN 配置的类，定义于 config.py 中。

以下是该类的一些属性说明：

• output_dim：GAN 的输出分辨率（128、256 或 512），适用于预训练模型；
• z_dim：噪声向量的维度（128，适用于预训练模型）；
• class_embed_dim：类别嵌入向量的维度（128，适用于预训练模型）；
• channel_width：每个通道的宽度（128，适用于预训练模型）；
• num_classes：训练数据集中的类别数量，例如 ImageNet（1000 类，适用于预训练模型）；
• layers：层定义列表。每项定义由三元组组成：[是否为上采样层？（布尔值）、输入通道数（整数）、输出通道数（整数）]；
• attention_layer_position：自注意力层在层结构中的位置（适用于预训练模型，为第 8 层）；
• eps：用于谱归一化和批量归一化层的 ε 值（适用于预训练模型，为 1e-4）；
• n_stats：与不同截断值（0 到 1）相关的批量归一化层的预计算统计数量（适用于预训练模型，为 51）。

模型

BigGAN 是一个 PyTorch 模型（torch.nn.Module），其定义位于 model.py 中。该模型由类别嵌入（一个线性层）和包含一系列卷积及条件批归一化的生成器组成。目前尚未实现判别器，因为其预训练权重尚未公开。

输入和输出与 TensorFlow 模型的输入和输出完全一致。

我们在此详细说明它们。

BigGAN 的 输入 包括：

• z：形状为 [batch_size, config.z_dim] 的 torch.FloatTensor，其中噪声是从截断正态分布中采样的；
• class_label：可选的 torch.LongTensor，形状为 [batch_size, sequence_length]，包含从 [0, 1] 中选择的标记类型索引。类型 0 对应于“句子 A”，类型 1 对应于“句子 B”标记（更多详情请参阅 BERT 论文）；
• truncation：一个介于 0（不包括）和 1 之间的浮点数，用于控制生成噪声向量时所用截断正态分布的截断程度。此截断值用于在预计算好的批归一化层统计量（均值和方差）集合中进行选择。

BigGAN 的 输出 是一个形状为 [batch_size, 3, resolution, resolution] 的张量，其中分辨率根据模型的不同可以是 128、256 或 512。

工具：图像、噪声、ImageNet 类别

我们提供了一些用于使用该模型的实用方法，这些方法定义在 utils.py 中。

以下是这些方法的详细信息：

• truncated_noise_sample(batch_size=1, dim_z=128, truncation=1., seed=None)：

  创建一个截断噪声向量。

  • 参数：
    • batch_size：批次大小。
    • dim_z：z 的维度。
    • truncation：使用的截断值。
    • seed：随机数生成器的种子。
  • 输出： 形状为 (batch_size, dim_z) 的数组。

• convert_to_images(obj)：

  将 BigGAN 的输出张量转换为图像列表。

  • 参数：
    • obj：形状为 (batch_size, channels, height, width) 的张量或 NumPy 数组。
  • 输出：
    • 大小为 (height, width) 的 Pillow 图像列表。

• save_as_images(obj, file_name='output')：

  将 BigGAN 的输出张量转换并保存为一系列图像文件。

  • 参数：
    • obj：形状为 (batch_size, channels, height, width) 的张量或 NumPy 数组。
    • file_name：保存路径及文件名前缀。图像将被保存为 file_name_{image_number}.png。

• display_in_terminal(obj)：

  将 BigGAN 的输出张量转换并在终端中显示。此函数使用 libsixel 库，仅能在支持 libsixel 的终端中运行。更多信息请参阅 https://github.com/saitoha/libsixel。

  • 参数：
    • obj：形状为 (batch_size, channels, height, width) 的张量或 NumPy 数组。
    • file_name：保存路径及文件名前缀。图像将被保存为 file_name_{image_number}.png。

• one_hot_from_int(int_or_list, batch_size=1)：

  根据类别索引或类别索引列表创建独热编码向量。

  • 参数：
    • int_or_list：ImageNet 类别的整数，或整数列表（范围在 0 到 999 之间）。
    • batch_size：批次大小。
      • 如果 int_or_list 是一个整数，则生成一批相同的类别。
      • 如果 int_or_list 是一个列表，则应满足 len(int_or_list) == batch_size。
  • 输出： 形状为 (batch_size, 1000) 的数组。

• one_hot_from_names(class_name, batch_size=1)：

  根据 ImageNet 类别的名称（如“网球”、“雏菊”等）创建独热编码向量。我们使用 NLTK 的 WordNet 搜索来尝试找到对应的 ImageNet 同义词集，并取第一个结果。如果无法直接找到，则会查找该类别的下位词和上位词。

  • 参数：
    • class_name：包含 ImageNet 物体名称的字符串。
  • 输出： 形状为 (batch_size, 1000) 的数组。

下载和转换脚本

用于从 TensorFlow Hub 下载并转换 TensorFlow 模型的脚本位于 ./scripts 中。

这些脚本可以直接使用，命令如下：

./scripts/download_tf_hub_models.sh
./scripts/convert_tf_hub_models.sh

PyTorch Pretrained BigGAN 快速上手指南

本指南帮助开发者快速在 PyTorch 环境中使用 DeepMind 的 BigGAN 预训练模型，生成高分辨率自然图像。

环境准备

• 操作系统：Linux / macOS / Windows
• Python 版本：推荐 Python 3.6+
• 核心依赖：
  • PyTorch >= 1.0.1
  • torchvision
• 可选依赖（如需使用模型转换脚本或 ImageNet 工具）：
  • TensorFlow
  • NLTK
  • libsixel（用于终端显示图像）

💡 提示：国内用户可使用清华或阿里镜像源加速 PyTorch 安装：

pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

方式一：仅使用生成模型（推荐）

若只需加载预训练模型进行图像生成，直接通过 pip 安装即可：

pip install pytorch-pretrained-biggan

方式二：完整功能安装（含转换脚本与工具）

如需下载并转换 TensorFlow Hub 上的原始模型，或使用 ImageNet 类别名称查询等功能，请克隆仓库并安装全部依赖：

git clone https://github.com/huggingface/pytorch-pretrained-BigGAN.git
cd pytorch-pretrained-BigGAN
pip install -r full_requirements.txt

💡 国内用户可考虑使用 Gitee 镜像加速克隆（如有）：

git clone https://gitee.com/mirror/pytorch-pretrained-BigGAN.git

基本使用

以下是最小可用示例，加载 biggan-deep-256 模型并生成三张图像（肥皂泡、咖啡、蘑菇）：

import torch
from pytorch_pretrained_biggan import (BigGAN, one_hot_from_names, truncated_noise_sample,
                                       save_as_images)

# 可选：开启日志查看详细加载过程
import logging
logging.basicConfig(level=logging.INFO)

# 加载预训练模型（自动下载并缓存）
model = BigGAN.from_pretrained('biggan-deep-256')

# 设置截断值（控制多样性与质量平衡）
truncation = 0.4

# 准备类别向量（支持英文 ImageNet 类别名）
class_vector = one_hot_from_names(['soap bubble', 'coffee', 'mushroom'], batch_size=3)

# 生成截断噪声向量
noise_vector = truncated_noise_sample(truncation=truncation, batch_size=3)

# 转为 Tensor 并移至 GPU（如有）
noise_vector = torch.from_numpy(noise_vector).to('cuda' if torch.cuda.is_available() else 'cpu')
class_vector = torch.from_numpy(class_vector).to('cuda' if torch.cuda.is_available() else 'cpu')
model.to('cuda' if torch.cuda.is_available() else 'cpu')

# 生成图像
with torch.no_grad():
    output = model(noise_vector, class_vector, truncation)

# 保存为 PNG 图片
save_as_images(output.to('cpu'), file_name='my_biggen_output')

运行后将生成 my_biggen_output_0.png、my_biggen_output_1.png、my_biggen_output_2.png 三张 256x256 分辨率的高质量图像。

📌 可用类别名参考 ImageNet 数据集，如 'tennis ball', 'daisy', 'goldfish' 等。