最近,我们发布了diffusion-deploy,以加快在GPU(~4倍速度,基于TensorRT)和CPU(~12倍速度,基于IntelOpenVINO)上的推理。
与这个 repo 集成后,small-stable-diffusion可以在CPU上仅用10秒生成图像。
这里我们给出了一个small-stable-diffusion, 和stablte diffusion相似的图片质量,但是模型小了近一半
一些样例:
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'})
这个模型是由stable-diffusion v1-4初始化的。由于模型结构与stable-diffusion不一样,参数数量较少,不能直接利用stable-diffusion的参数。因此,small-stable-diffusion设置layers_per_block=1
,并选择原始stable-diffusion中每个块的第一层来初始化小模型。
初始化后,模型在8xA100 GPUS中被训练了1100k步。训练过程由三个阶段组成。第一阶段是一个简单的预训练阶段。在后两个阶段,利用原始的stable-diffusion作为教师模型来蒸馏小模型。在所有阶段,只有unet中的参数被训练,其他参数被冻结。
硬件: 8 x A100-80GB GPUs
优化器: AdamW
阶段 1 - 预训练模型的unet部分。
阶段 2 - 使用stable-diffusionv1-4作为教师来提炼模型。除了地面真实,这个阶段的训练还使用教师模型产生的软标签(pred_noise
)
阶段 3 - 使用stable-diffusionv1-5作为老师来蒸馏模型。 用到了Knowledge Distillation of Transformer-based Language Models Revisited
中的一些技术,包括除了软标签之外的基于相似性的层匹配。
我们使用以下数据集来训练模型。
@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}
}