重写深度生成模型

在本文中，我们探讨是否可以通过让用户直接更改权重，而不是使用数据集进行训练，来重新编程深度网络以遵循不同的规则。

直接重写StyleGANv2的权重，将马匹重新编程为戴帽子的样子。

什么是模型重写？

我们提出了模型重写的任务，其目标是添加、移除以及改变预训练深度网络的语义和物理规则。虽然现代图像编辑工具通过操作单个输入图像来实现用户指定的目标，但我们允许用户通过编辑生成模型使其遵循修改后的规则，从而合成无限数量的新图像。

为什么需要重写模型？

希望直接重写深度网络的原因有两个：

1. 为了深入了解深度网络如何组织其知识。
2. 为了让富有创造力的用户能够快速创建没有现有数据集支持的新模型。

模型重写设想了一种根据用户意图构建深度网络的方式。与其将网络局限于模仿我们已有的数据，不如通过重写使深度网络能够模拟一个遵循用户期望新规则的世界。

重写深度生成模型。
大卫·鲍、史蒂文·刘、王通舟、朱俊彦、安东尼奥·托拉尔巴。
ECCV 2020（口头报告）。
MIT CSAIL与Adobe Research。

2分钟 视频预览

ECCV 2020 预印本

ECCV 2020 演讲视频

演示 幻灯片

网站

演示Colab 笔记本

我们的方法通过重写生成器的权重来改变生成规则。不同于对单张图像进行编辑，我们的方法直接编辑生成器，因此可以利用修改后的规则潜在地合成和操控无限多的图像。规则可以通过多种方式改变，例如移除水印等图案、添加人物等对象，或替换定义，比如让树木从塔楼上生长出来。

我们的方法基于这样一个假设：生成器的权重充当线性联想记忆。每一层存储着键与值之间的映射关系，其中键表示有意义的上下文，而值则决定输出。

示例结果

以下结果展示了对StyleGANv2中单一规则的修改。在每种情况下， 由用户选择的四个示例（顶部行中央）确立了要被重写的规则的上下文， 而“复制并粘贴”的示例（顶部行左右两侧）则表明了用户希望如何改变模型。

下图中的网格显示了成对的输出：对于每一对，第一个是原始未修改的StyleGANv2的输出， 第二个则是应用了我们方法后，按照用户的意图修改过的StyleGANv2的输出。

---

首先：将儿童眉毛的规则改为看起来像浓密的小胡子。

---

将尖顶塔楼的规则改为让它们长出树木。

---

将马头顶部的规则改为给马戴上帽子。

---

将皱眉的表情改为微笑。

---

通过改变绘制空白墙的规则，移除建筑物中的主窗户。

使用提示

代码使用PyTorch运行。

• 方法和界面位于/rewrite
• 笔记本位于/notebooks：请参阅rewriting-interface.ipynb以获取演示UI。
• 定量实验位于/metrics，剖析工具位于/utils。 第5.1节的实验可以通过/experiments.sh复现， 而第5.2节的实验可以通过/watermarks.sh复现。第5.3节的实验则在/notebooks/reflection-rule-change.ipynb笔记本中运行。

环境设置

该代码设计用于在Python 3.6上运行最新版本的PyTorch (1.4+)， 并搭配CUDA 10.1和cuDNN 7.6.0。运行setup/setup_renv.sh即可创建包含所需依赖项的conda环境。

重写您自己的模型

要编辑您自己的模型，请按以下步骤操作：

• 使用标准的StyleGAN v2 TensorFlow发布版训练您的GAN，该版本可在https://github.com/NVlabs/stylegan2找到。
• 使用rosinality提供的convert_weight.py工具将您训练好的权重转换为PyTorch格式， 该工具可在https://github.com/rosinality/stylegan2-pytorch中找到。
• 随后，这些权重文件可以直接加载到我们序列化的StyleGAN移植版本中，如下所示：

# 分辨率（size）和风格维度（style_dim和n_mlp）是您训练时所使用的架构尺寸。
# 如果需要，此处也可以应用截断技巧（如果不使用，则设为truncation=1.0）。
# 注意，mconv='seq'会将优化后的调制卷积拆分为独立的操作，
# 这样重写器就可以直接检查底层的卷积。
model = SeqStyleGAN2(size=256, style_dim=512, n_mlp=8, truncation=0.5, mconv='seq')

# 加载指数移动平均模型权重，并将其放到 GPU 上。
state_dict = torch.load('your_model.pt')
model.load_state_dict(state_dict['g_ema'], latent_avg=state_dict['latent_avg'])
model.cuda()

• 创建一个 ganrewrite.SeqStyleGanRewriter 实例来编辑你的模型

layernum = 8 # 或者你想要修改的任意层
sample_size = 1000 # 用于计算统计量的一小批图像
zds = zdataset.z_dataset_for_model(model, size=sample_size)
gw = SeqStyleGanRewriter(
    model, zds, layernum,
    cachedir='experiments')

• 最后，要指定并创建修改，你可以使用 rewriteapp.GanRewriteApp 界面（假设你在 Jupyter Notebook 中运行）。该界面可用于尝试编辑，并保存和加载包含已存储模型编辑规范的 JSON 文件。

savedir = 'masks'
interface = rewriteapp.GanRewriteApp(gw, size=256, mask_dir=savedir, num_canvases=32)

• 若要批量生成经过编辑后的模型输出的图像（无需 Jupyter Notebook 界面），可以执行以下操作，然后像往常一样对修改后的模型输出进行采样。示例请参见 metrics/sample_edited.py。

saved_edit = 'masks/my_edit.json'
gw.apply_edit(json.load(saved_edit), rank=1)

Rewriting 快速上手指南

Rewriting 是一个用于直接修改预训练深度生成模型（如 StyleGANv2）权重的工具。它允许用户通过编辑模型规则来添加、移除或改变生成图像的语义特征（例如给马加上帽子、将塔顶变成树），而无需重新训练模型或准备新数据集。

环境准备

本项目基于 PyTorch 开发，推荐在 Linux 环境下运行。

• 操作系统: Linux
• Python 版本: 3.6+
• PyTorch 版本: 1.4+
• CUDA: 10.1
• cuDNN: 7.6.0
• 依赖管理: Conda (推荐)

注意：国内用户建议在创建 Conda 环境时配置清华或中科大镜像源以加速依赖下载。

安装步骤

1. 克隆项目代码

   git clone https://github.com/davidbau/rewriting.git
   cd rewriting
   

2. 创建 Conda 环境并安装依赖 项目提供了自动化脚本 setup_renv.sh 来配置所需的环境。

   bash setup/setup_renv.sh
   

   如果脚本执行失败，请手动激活环境并安装 requirements.txt 中的依赖。

3. 验证安装 激活环境后，尝试运行示例 Notebook 或检查 PyTorch 是否识别到 GPU：

   conda activate rewriting_env  # 环境名称可能因脚本而异，请根据实际输出调整
   python -c "import torch; print(torch.cuda.is_available())"
   

基本使用

本工具的核心流程是：加载预训练的 StyleGANv2 模型 -> 初始化重写器 -> 使用交互界面或代码应用编辑规则。

1. 准备模型权重

你需要一个标准的 StyleGANv2 模型（.pt 格式）。如果你只有 TensorFlow 版本的权重，请先使用 rosinality/stylegan2-pytorch 中的 convert_weight.py 进行转换。

2. 加载模型与初始化重写器

以下代码展示了如何加载模型并初始化重写接口（建议在 Jupyter Notebook 中运行）：

import torch
from stylegan2 import SeqStyleGAN2  # 假设已导入项目内的序列化模型定义
import zdataset
from ganrewrite import SeqStyleGanRewriter
from rewriteapp import GanRewriteApp

# 1. 初始化模型
# 参数需与你训练时的架构一致 (size, style_dim, n_mlp)
# mconv='seq' 是关键，它将卷积操作拆分以便重写器访问底层卷积
model = SeqStyleGAN2(size=256, style_dim=512, n_mlp=8, truncation=0.5, mconv='seq')

# 2. 加载权重
# 加载 EMA (指数移动平均) 权重并移至 GPU
state_dict = torch.load('your_model.pt')
model.load_state_dict(state_dict['g_ema'], latent_avg=state_dict['latent_avg'])
model.cuda()

# 3. 初始化重写器
layernum = 8          # 指定要修改的网络层数
sample_size = 1000    # 用于计算统计信息的小样本图像数量
zds = zdataset.z_dataset_for_model(model, size=sample_size)

gw = SeqStyleGanRewriter(
    model, zds, layernum,
    cachedir='experiments'
)

# 4. 启动交互界面 (仅在 Notebook 中有效)
# 该界面允许你选择源区域和目标区域来定义规则变化
savedir = 'masks'
interface = GanRewriteApp(gw, size=256, mask_dir=savedir, num_canvases=32)

3. 应用编辑并生成图像

在交互界面中完成规则定义并保存为 JSON 文件后，可以通过代码批量应用编辑并生成新图像：

import json

# 加载之前保存的编辑规则
saved_edit = 'masks/my_edit.json'

# 应用编辑到模型 (rank=1 表示应用主要编辑方向)
gw.apply_edit(json.load(open(saved_edit)), rank=1)

# 此时 model 已被修改，可像往常一样采样生成图像
# 参考 metrics/sample_edited.py 查看完整的生成示例

快速体验

如果你想直接体验效果而无需配置本地环境，可以使用官方提供的 Colab 演示：

• Demo Colab Notebook: https://colab.research.google.com/github/davidbau/rewriting/blob/master/notebooks/rewriting-interface.ipynb