空间增量生成引擎 (SIGE)

论文 | 项目 | 幻灯片 | YouTube

[最新消息!] SIGE已被T-PAMI接收！

[最新消息!] SIGE现已支持Stable Diffusion和Mac MPS后端！我们还发布了DDPM在M1 Macbook Pro上的交互式演示代码！

[最新消息!] SIGE已被NeurIPS 2022接收！我们的代码和基准数据集已公开可用！

预览图我们提出了空间增量生成引擎（SIGE），这是一个针对图像编辑应用选择性地在编辑区域执行计算的引擎。计算量和延迟是针对单次前向传播测量的。对于上述示例，SIGE显著降低了SDEdit与DDPM（4-6倍）、Stable Diffusion（8倍）和GauGAN（15倍）的计算量，同时保持图像质量。当与现有的模型压缩方法（如GAN压缩）结合时，它进一步将GauGAN的计算量减少了47倍。在NVIDIA RTX 3090上，SIGE实现了高达7.2倍的加速。

条件GAN和扩散模型的高效空间稀疏推理
李牧阳、林吉、孟晨林、Stefano Ermon、韩松和朱俊彦
卡内基梅隆大学、麻省理工学院和斯坦福大学
NeurIPS 2022

演示

相比原始DDPM，SIGE在M1 MacBook Pro GPU上实现了2倍更少的转换时间，因为我们选择性地在编辑区域执行计算。

概述

基于分块的稀疏卷积概述。对于网络中的每个卷积F_l，我们将其封装为SIGE Conv_l。原始图像的激活已经预先计算好。在获得编辑后的图像时，我们首先计算原始图像和编辑后图像之间的差异掩码，并将掩码缩减为活跃块索引以定位编辑区域。在每个SIGE Conv_l中，我们根据缩减后的索引直接从编辑后的激活A_l^edited中收集活跃块，沿批次维度堆叠这些块，并将它们输入F_l。如果F_l是步长为1的3×3卷积，收集的块会有宽度为2的重叠。在从F_l获得输出块后，我们将它们散布回F_l(A_l^original)以获得编辑后的输出，这近似于F_l(A_l^edited)。

性能

效率

在1.2%的编辑下，SIGE可以将DDPM、Progressive Distillation和GauGAN的计算量减少7-18倍，在NVIDIA RTX 3090上实现2-4倍加速，在Apple M1 Pro GPU上实现3-5倍加速，在M1 Pro CPU上实现4-14倍加速。当与GAN压缩结合时，它进一步将GauGAN的计算量减少50倍，在M1 Pro CPU上实现38倍加速。请查看我们的论文以获取更多细节和结果。

质量

不同编辑大小下的定性结果。PD是Progressive Distillation。我们的方法很好地保持了原始模型的视觉保真度，而不会丢失全局上下文。 quality-stable-diffusion

在NVIDIA RTX 3090上测量的Stable Diffusion在图像修复和编辑方面的更多定性结果。

参考文献：

去噪扩散概率模型（DDPM），Ho等人，ICLR 2020
去噪扩散隐式模型（DDIM），Song等人，ICLR 2021
用于扩散模型快速采样的渐进蒸馏，Salimans等人，ICLR 2022
具有空间自适应归一化的语义图像合成（GauGAN），Park等人，CVPR 2019
GAN压缩：交互式条件GAN的高效架构，Li等人，CVPR 2020
使用潜在扩散模型的高分辨率图像合成，Rombach等人，CVPR 2022

先决条件

Python3
CPU、M1 GPU或NVIDIA GPU + CUDA CuDNN
PyTorch >= 1.7。对于M1 GPU支持，请安装PyTorch>=2.0。

入门指南

安装

安装PyTorch后，您应该能够通过PyPI安装SIGE

pip install sige

或通过GitHub：

pip install git+https://github.com/lmxyy/sige.git

或本地安装以进行开发

git clone git@github.com:lmxyy/sige.git
cd sige
pip install -e .

对于MPS后端，请设置环境变量：

export PYTORCH_ENABLE_MPS_FALLBACK=1

使用示例

请参阅example.py以获取最小SIGE卷积示例。首先按照上述说明安装SIGE，并使用以下命令安装torchprofile：

pip install torchprofile

然后您可以运行：

python example.py

我们还提供了示例。

基准测试

要重现DDPM和Progressive Distillation的结果，或下载LSUN Church编辑数据集，请按照diffusion/README.md中的说明操作。

要重现GauGAN和GAN Compression的结果，或下载Cityscapes编辑数据集，请按照gaugan/README.md中的说明操作。

引用

如果您在研究中使用此代码，请引用我们的论文。

@inproceedings{li2022efficient,
  title={Efficient Spatially Sparse Inference for Conditional GANs and Diffusion Models},
  author={Li, Muyang and Lin, Ji and Meng, Chenlin and Ermon, Stefano and Han, Song and Zhu, Jun-Yan},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2022}
}

致谢

我们的代码基于SDEdit、ddim、diffusion_distillation、gan-compression、dpm-solver和stable-diffusion开发。我们参考了sbnet实现基于平铺的稀疏卷积算法。我们的工作还受到torchsparse中收集/分散实现的启发。

我们感谢torchprofile用于MACs测量，clean-fid用于FID计算，以及drn用于Cityscapes mIoU计算。

我们感谢丁耀耀、叶子豪、郑连敏、唐浩天和朱立庚对引擎设计的有益评论。我们还感谢George Cazenavette、邓康乐、高睿涵、陆道涵、王圣宇和张炳良的宝贵反馈。该项目部分得到NSF、MIT-IBM Watson AI Lab、快手公司和索尼公司的支持。