small-stable-diffusion-v0
  • 模型资讯
  • 模型资料

Small Stable Diffusion V0

最近,我们发布了diffusion-deploy,以加快在GPU(~4倍速度,基于TensorRT)和CPU(~12倍速度,基于IntelOpenVINO)上的推理。
与这个 repo 集成后,small-stable-diffusion可以在CPU上仅用10秒生成图像
这里我们给出了一个small-stable-diffusion, 和stablte diffusion相似的图片质量,但是模型小了近一半
一些样例:

More samples.

代码使用范例


import torch
import cv2
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

task = Tasks.text_to_image_synthesis
model_id = 'modelscope/small-stable-diffusion-v0'
# 基础调用
pipe = pipeline(task=task, model=model_id, model_revision='v1.0.2')
output = pipe({'text': 'an apple'})

Training

Initialization

这个模型是由stable-diffusion v1-4初始化的。由于模型结构与stable-diffusion不一样,参数数量较少,不能直接利用stable-diffusion的参数。因此,small-stable-diffusion设置layers_per_block=1,并选择原始stable-diffusion中每个块的第一层来初始化小模型。

Training Procedure

初始化后,模型在8xA100 GPUS中被训练了1100k步。训练过程由三个阶段组成。第一阶段是一个简单的预训练阶段。在后两个阶段,利用原始的stable-diffusion作为教师模型来蒸馏小模型。在所有阶段,只有unet中的参数被训练,其他参数被冻结。

  • 硬件: 8 x A100-80GB GPUs

  • 优化器: AdamW

  • 阶段 1 - 预训练模型的unet部分。

    • Steps: 500,000
    • Batch: batch size=8, GPUs=8, Gradient Accumulations=2. Total batch size=128
    • Learning rate: 预热10,000步到1e-5,然后保持稳定
  • 阶段 2 - 使用stable-diffusionv1-4作为教师来提炼模型。除了地面真实,这个阶段的训练还使用教师模型产生的软标签(pred_noise

    • Steps: 400,000
    • Batch: batch size=8, GPUs=8, Gradient Accumulations=2. Total batch size=128
    • Learning rate: 预热5,000步到1e-5,然后保持稳定
    • Soft label weight: 0.5
    • Hard label weight: 0.5
  • 阶段 3 - 使用stable-diffusionv1-5作为老师来蒸馏模型。 用到了Knowledge Distillation of Transformer-based Language Models Revisited中的一些技术,包括除了软标签之外的基于相似性的层匹配。

    • Steps: 200,000
    • Batch: batch size=8, GPUs=8, Gradient Accumulations=2. Total batch size=128
    • Learning rate: 预热5,000步到1e-5,然后保持稳定
    • Softlabel weight: 0.5
    • Hard label weight: 0.5

Training Data

我们使用以下数据集来训练模型。

  1. LAION-2B en aesthetic
  2. LAION-Art
  3. LAION-HD

Citation

@article{Lu2022KnowledgeDO,
  title={Knowledge Distillation of Transformer-based Language Models Revisited},
  author={Chengqiang Lu and Jianwei Zhang and Yunfei Chu and Zhengyu Chen and Jingren Zhou and Fei Wu and Haiqing Chen and Hongxia Yang},
  journal={ArXiv},
  year={2022},
  volume={abs/2206.14366}
}