关于 Stable Diffusion 模型优化分享

一、背景

随着技术的飞速发展，人工智能技术已经成为推动社会变革的关键力量。在这个充满创新的时代，oneAPI 技术堆栈崭露头角，为构建各种创新解决方案提供了巨大的潜力。在这一背景下，本次竞赛旨在深入探讨人工智能技术在特定领域的应用，涵盖机器学习、深度学习和数据分析等多个方面，同时为参赛者提供实践机会，通过解决问题和实现功能，更好地理解和运用 oneAPI 技术。

Stable Diffusion 是 2022 年发布的深度学习图像化生成模型，它主要用于根据文本的描述产生详细图像，尽管它也可以应用于其他任务，如内补绘制、外补绘制，以及在提示词指导下产生图生图的翻译。

Stable Diffusion 技术作为一种先进的生成模型，具有在生成图像任务中表现出色的潜力。然而，在实际部署中，要确保模型在端侧设备上的高效运行，需要面对一系列挑战，包括性能瓶颈和资源利用率。通过模型优化方案，参赛者将深入挖掘 Stable Diffusion 技术的性能潜力，结合 oneAPI 技术堆栈，实现在指定硬件平台上的部署优化，为生成图任务提供更高效、更稳定的解决方案。本篇文章就我参与的比赛的一些心得感受，优化思路作为分享内容呈现给大家，这和上一篇不同，是一个全新的优化方向，本人也在比赛中实现了部分内容，话不多说，现就就开始今天的分享！

二、原理解读

文生图任务是指将一段文本输入到 SD 模型中，经过一定的迭代次数，SD 模型输出一张符合输入文本描述的图片。

该模型主要可以分为三个部分：

1. 变分编码器 Vector Quantised Variational AutoEncoder，VQ-VAE

2. 扩散模型 Diffusion Model， DM

3. 条件控制器 Conditioning

其中主要的 VAE 由编码器（Encoder）和解码器（Decoder）两部分组成：

Step 1. 输入图片 Input 通过编码器被到转换到潜在空间，得到潜在空间的图片表示 Latent Image Input

Step 2. 配合 Conditioning，Diffusion Model 对 Latent Image Input 进行处理，产生 Latent Image Output

Step 3. 解码器将由 Diffusion Model 产生的 Latent Image Output 映射回像素空间，得到输出图片 Output

图片通过 VAE 转换到低维空间，配合 Conditioning 的 DM 产生新的变量，再通过 VAE 将生成的变量转换为图片。

例如赛题要求：

1. Prompt 输入："a photo of an astronaut riding a horse on mars" 

2. Negative Prompt 输入："low resolution, blurry" 

图片输出：512*512，24 Bit，PNG 格式

利用 VAE 的编码器将输入图片 Input 降维，得到 Latent Image Input。利用训练好的 DM，不断对图片进行噪声预测，并对 Latent Image Input 进行去噪，经过一定步骤后得到去除了噪声的 Latent Image Output，最终通过 VAE 的解码器得到输出图片 Output。

通过 text prompt 得到的 Embedding 暂时还无法直接使用，还需要通过 Transfomer 进行再加工才能喂给属于 DM 的噪声预测器。值得一提的是，Transformer 是 SD 能够支持多模态的重要原因，它不仅能够处理 text prompt 生成的 embedding，还能够处理类似图片、深度图等输入，将其转化为噪声预测器能够使用的数据。

Transfomer 的输出会被噪声预测器多次利用，并且由于 Transfomer 的 Cross Attention 机制，它能够正确的利用 text prompt 中的内容。并且由于 Transfomer 的 Self Attention 机制，prompt 能够被正确解读，例如"a photo of an astronaut riding a horse on mars"，SD 会将"astronaut"和"horse"组合。然后利用这个信息去影响噪声预测器的输出，让 DM 的逆向过程朝着带有”an astronaut riding a horse“的图像生成。

二、优化方向解析

我的项目具体实现是致力于解决在文生成图任务中，模型规模庞大导致的高存储需求和计算开销大的问题。具体方法是通过采用渐进式模型剪枝与量化策略配和 CPU 与 GPU 的混合使用，能够在不损失生成质量的前提下，逐步减小模型的大小，并提高模型的推理速度。（项目中具体实现了模型的剪枝，由于时间问题和学习学校课程，文章中的其他优化点还没有具体实现，其它的优化方案在本文中简要的做了说明，有兴趣的小伙伴可以联系 wx 一起探讨实现）

在模型优化方面，我主要关注神经元剪枝算法，通过精细的剪枝策略降低了模型的冗余部分，同时利用 OpenVINO 工具对模型进行文生成图预处理。利用 OpenVINO 工具套件的 Layout API 对输入进行预处理，一点一点微调，我在不牺牲生成质量和大小的前提下，逐步减小模型大小并提高推理速度，一点点实现了异步执行与 Pipeline 并行性，充分发挥此次大赛提供的硬件资源的优势，为端到端性能提升和硬件适应性提供了一体化的解决方案。（大家有更好的优化方案、想法可以一起讨论）

三、模型压缩方案

OpenVINO 工具套件提供了一系列的模型优化工具，包括模型剪枝和量化等等，我的思路和实现也主要是基于这两套工具来对 SD 模型在比赛提供的硬件上进行模型训练和优化的。具体如下：

首先必不可少的当然是对于工具的利用：在改进数据管道和预处理加速中， OpenVINO 工具中有许多可用的方法 。

3.0 利用工具优化：

我们的实例中主要运用的是以下两种 API 对于示例模型进行了一个简单优化处理，在 GPU 占用率上有了明显可见的下降，且相同配置参数下文生图时间间隔下降了 0.4s。具体实现思路和伪代码如下（主要用于大家学习思路）：

一、使用 OpenVINO 对输入进行预处理：

1. 在对 Static Diffusion 模型进行优化时，采用 OpenVINO 进行输入预处理是至关重要的。以下是将预处理步骤集成到模型中的具体思路：

   声明 Tensor 格式： 首先，从实际用户数据中声明模型输入的 Tensor 格式，包括形状、布局、精度、颜色格式等。这样的声明有助于确保模型输入与实际推理数据的格式相匹配。

   描述预处理步骤： 确定需要应用于用户数据的预处理步骤序列。这可能包括均值调整、尺度缩放、通道反转等，以确保输入数据满足模型的要求。通过 OpenVINO 的模型转换 API，可以方便地描述和配置这些预处理步骤。

   指定模型数据格式： 对于 Static Diffusion 模型，模型的精度和形状通常是已知的，但需要指定其他信息，如布局等。通过 OpenVINO 提供的模型数据格式参数，可以将模型与实际推理数据正确对齐。

   集成到模型中： 完成预处理步骤后，将这些步骤集成到模型中。通过 OpenVINO 提供的模型构建功能，可以轻松构建具备预处理功能的模型。

   python # 示例代码 from openvino.tools.preprocessor import Preprocessor ​ # 1. 声明 Tensor 格式 tensor_format = {...} # 根据实际情况填写 ​ # 2. 描述预处理步骤 preprocess_steps = [...] # 根据实际需求填写 ​ # 3. 指定模型数据格式 model_format = {...} # 根据实际情况填写 ​ # 4. 集成到模型中 preprocessor = Preprocessor(tensor_format, preprocess_steps, model_format) print(f'Dump preprocessor: {preprocessor}') model = preprocessor.build()

二、具体使用 OpenVINO Layout API

1. 定义输入和输出的 Tensor 格式： 利用 Layout API 声明 Static Diffusion 模型输入和输出的 Tensor 格式，包括形状、布局、精度等信息。这有助于确保模型的输入和输出与实际推理数据的格式相匹配。

pythonfrom openvino.runtime import Layoutinput_layout = Layout('NCHW')output_layout = Layout('NCHW')

1. 进行模型修改和预处理： 应用 Layout API 中的布局信息，执行与模型修改相关的操作，例如应用预处理步骤、调整图像大小等。这确保了输入数据在推理前得到正确的处理，以适应模型的期望。

python# 应用预处理步骤preprocessor = Preprocessor(input_layout, preprocess_steps, output_layout)model = preprocessor.build()

1. 设置 Batch 大小： 利用 Layout API 的预定义名称，设置模型的 Batch 大小，以便更好地处理多个输入数据。

pythonfrom openvino.runtime import layout_helpersbatch_idx = layout_helpers.batch_idx(input_layout)model.set_batch_size(batch_size)

1. 提高模型输入和输出的可读性： 利用 Layout API 中的布局信息，提高模型输入和输出的可读性，让用户更容易理解各个维度代表的含义。

pythonprint(f'Model Input Layout: {input_layout}')    # 输出 [N,C,H,W]print(f'Model Output Layout: {output_layout}')   # 输出 [N,C,H,W]

可以在 Static Diffusion 模型中更好地理解和处理输入输出数据，降低 pipeline 端到端延迟。

3.1 模型剪枝

模型剪枝（Pruning）也叫模型稀疏化，不同于模型量化对每一个权重参数进行压缩，稀疏化方法是尝试直接“删除”部分权重参数。模型剪枝的原理是通过剔除模型中 “不重要” 的权重，使得模型减少参数量和计算量，同时尽量保证模型的精度不受影响。

我们的主要实现方式是利用 OpenVINO 工具套件的模型剪枝和量化功能，有选择性地减小模型的规模，去除冗余参数，以适应端侧设备的资源限制。然后，借助 OpenVINO 的量化功能，将模型参数映射到低精度表示，从而显著减小模型的体积。降低模型在存储和传输中的开销，同时提高了推理速度，使得整个 Pipeline 的性能和端到端性能提升进一步优化。

我们在优化 Static Diffusion 模型时，也可以结合权重稀疏的思路，可通过度量权重的绝对值大小，定义阈值进行稀疏化，并动态调整阈值以适应模型动态变化。

通过模型剪枝技术实现，保留对模型影响较大的权重，去除对模型影响较小的权重，以达到模型的精简和加速。在稀疏化后需进行重新训练或微调以维持模型性能，并全面评估 SD 模型文生图推理速度、内存、GPU、CPU 占用和生成图像质量等指标，渐进式的进行动态去除，在生成任务中定一个计时器代码段，用于计算前后剪枝处理的时间比，计算出最优的剪枝结果，从而获取最高的效率。（优化效果不显著）

之后对于这种优化方式已不抱多大的希望，但在我翻阅大量的论文和材料后发现一种神经元剪枝算法（由于时间限制，代码中未实现...）。我在这里简单介绍一下这种算法的实现思路：

在 Static Diffusion 模型中，首先定义 Average Percentage of Zeros（APoZ）指标，用于衡量神经元零激活的百分比。根据给定的公式，计算每个神经元的 APoZ 值，表示其对模型输出的影响程度。

制定神经元剪枝策略，基于计算得到的 APoZ 指标，设定阈值，将 APoZ 值低于阈值的神经元标记为不重要的神经元，进而进行剪枝操作。这一步骤可以通过权重置零或者神经元删除的方式实现。

由于不同层次的神经元对模型的影响程度不同，可以根据具体网络结构，对不同层次的神经元采用不同的剪枝策略。例如，在 CONV 层和 FC 层分别调整剪枝阈值，以满足 SD 模型的优化需求。在进行神经元剪枝后，需要对剪枝后的模型进行验证，以确保剪枝操作不会显著降低模型性能。若验证结果满足要求，可以进行微调以进一步维持模型的准确性。

注：这种算法根本上优化可视度肯定比权重稀疏的效果明显（主要原因可以看看下面），有兴趣的可以一起研究试试。在神经元修剪后，修剪后的网络要使用修剪前的权重进行初始化。 在最后一步中，需要重新训练网络以加强剩余的神经元以增强修剪后网络的性能。

3.1.1 神经元剪枝

神经元剪枝是将某个/些神经元从网络中删除。这会使得神经网络架构的尺寸降低。

采用神经元剪枝（Neurons Pruning）的优势在于，它有效降低了 Static Diffusion 模型的神经网络架构尺寸，同时保持了密集计算的特性，包括 input feature map 和 dense kernel 计算。这种剪枝方法依赖于硬件设备，而此次大赛英特尔官方提供的硬件设备 very 给力，实现起来效果肯定也很佳。

对比传统的权重剪枝方法（效果细微），传统剪枝方法中，由于一个 kernel 中的元素存在 0 元素和非 0 元素，其中后者是有效的数据，而前者则是无效的。在 Memory 中，数据是连续存储的，而计算单元从 Memory 中取出的一个 block 中的 0 元素是无效的，从而降低了带宽利用率。这使得在相同硬件设备上权重剪枝的计算模式相对神经元剪枝更加复杂，且存在一定的计算效率损失。这也就是我肯定这种优化方案可行性的依据！！！

3.1.2 权重剪枝：

我们可以使用一些库和工具来实现。在这里，我将使用 TensorFlow Model Optimization Toolkit 来进行权重剪枝。

pip install tensorflowpip install tensorflow_model_optimization

然后，我们可以按照以下方式修改脚本：

pythonimport asynciofrom tensorflow import kerasfrom stable_diffusion_tf.stable_diffusion import StableDiffusionimport argparsefrom PIL import Imagefrom PIL.PngImagePlugin import PngInfoimport osfrom tensorflow_model_optimization.python.core.sparsity.keras import prune_low_magnitude​os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"os.environ["CUDA_VISIBLE_DEVICES"] = "-1"​async def generate_image_async(args):    if args.mixed_precision:        print("Using mixed precision.")        keras.mixed_precision.set_global_policy("mixed_float16")​    # Load your model    generator = StableDiffusion(img_height=args.height, img_width=args.width, jit_compile=False)    model = generator.get_model()​    # Apply weight pruning    pruned_model = prune_low_magnitude(model)    pruned_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])​    # Load existing weights to pruned model (assuming the model is already trained)    pruned_model.set_weights(model.get_weights())​    # Continue with the rest of the code...​    generated_image = generator.generate(        args.prompt,        negative_prompt=args.negative_prompt,        num_steps=args.steps,        unconditional_guidance_scale=args.scale,        temperature=1,        batch_size=1,        seed=args.seed,    )​    # Save the generated image with prompt information in PNG format    png_info = PngInfo()    png_info.add_text('prompt', args.prompt)    Image.fromarray(generated_image[0]).save(args.output, png_info=png_info)​    print(f"Saved generated image at {args.output}")​async def main():    args = parse_arguments()    await generate_image_async(args)​def parse_arguments():    parser = argparse.ArgumentParser(description="Generate images using Stable Diffusion model.")        # ... (unchanged)​    return parser.parse_args()​if __name__ == "__main__":    asyncio.run(main())

这里使用了 TensorFlow Model Optimization Toolkit 中的 prune_low_magnitude 函数来进行权重剪枝。

3.2 CPU 与 GPU 权重切换

首先，对 Stable Diffusion（SD）模型进行权重划分，将模型的不同部分或模块的权重进行分类。这可以根据模型结构、层级或其他相关因素进行，确保权重的划分具有合理性和可行性。

初始加载到 CPU： 将整个模型或划分后的部分模块的权重初始加载到 CPU。这一步骤可以在模型初始化阶段完成，确保 CPU 上具有完整的初始权重。

推理时动态加载到 GPU： 在进行推理时，根据需要动态加载相应部分的权重到 GPU。这可以根据模型的实际运行情况和推理需求来灵活选择加载哪些权重。推理开始前，只加载与当前推理相关的部分权重，以降低 GPU 内存占用。

定期检查模型的性能和推理需求，根据模型的动态变化情况，定期更新 GPU 上的权重。这有助于在模型训练后的不同阶段，以及在处理不同输入时，动态调整 GPU 上的权重。优化 AI 生图模型在端侧设备上的 Pipeline 性能，在保证生图效果的情况下，降低 pipeline 端到端延迟，降低 pipeline 峰值内存占用

算法：

import torch​class SDModel:    def __init__(self):        # 初始化整个模型权重加载到CPU        self.model_cpu = self.load_model_to_cpu()​        # GPU上的权重字典，用于存储已加载到GPU的权重        self.weights_on_gpu = {}​    def load_model_to_cpu(self):        # 加载整个模型到CPU的逻辑，可以根据实际需求进行修改        model_cpu = ...​        return model_cpu​    def load_weights_to_gpu(self, module_name):        # 模拟加载指定模块的权重到GPU的逻辑，可以根据实际需求进行修改        weights = self.model_cpu.get_weights(module_name)        self.weights_on_gpu[module_name] = torch.tensor(weights).cuda()​    def inference(self, input_data, gpu_module):        # 动态加载权重到GPU        self.load_weights_to_gpu(gpu_module)​        # 在GPU上进行推理        result = self.model_cpu.inference_on_gpu(input_data, self.weights_on_gpu[gpu_module])​        return result​# 示例使用sd_model = SDModel()​# 推理时动态加载指定模块权重到GPUresult = sd_model.inference(input_data, gpu_module='specific_module')

这个算法示例中，通过初始化整个模型加载到 CPU，然后在推理时根据需要动态加载指定模块的权重到 GPU。在实际实现中，可以根据 SD 模型结构和性能进行更复杂的权重加载策略。代码在模型具体示例。

四、总结

虽然这次大赛匆匆忙忙，但是对于对于人工智能技术的发展和应用有了更深了解，对 OneAPI 以及英特尔® AI 分析工具套件有了更熟悉的了解和使用。

总之，对于本次的优化还是很欠缺的，期待下次相遇，我也会抓住时间，打实基础，争取获奖！这次的实现和对模型优化并不多，主要是有思路，但是精力受限，也未找下合作伙伴，前前后后都是一个人在忙，也遇到了好多困难，好几次想要放弃，但是既然选择了，是好是坏还未知呢！坚持就对啦，在寒假会继续研究优化方案，有兴趣的小伙伴可以联系我一起！

以上就是本次大赛的总结分享，大模型对于科技和世界的改变近在眼前，技术永远在进步，我们要做的就是优化技术，适应时代大发展，科技创造未来！