Diffusers 文档
Pruna
加入 Hugging Face 社区
并获得增强的文档体验
开始使用
Pruna
Pruna 是一个模型优化框架，提供量化、剪枝、缓存、编译等多种优化方法，用于加速推理和减少内存使用。优化方法概述如下。
| 技术 | 描述 | 速度 | 内存 | 质量 |
|---|---|---|---|---|
批处理器 | 将多个输入组合在一起同时处理，提高计算效率并缩短处理时间。 | ✅ | ❌ | ➖ |
缓存器 | 存储计算的中间结果以加速后续操作。 | ✅ | ➖ | ➖ |
编译器 | 通过针对特定硬件的指令优化模型。 | ✅ | ➖ | ➖ |
蒸馏器 | 训练一个更小、更简单的模型来模仿一个更大、更复杂的模型。 | ✅ | ✅ | ❌ |
量化器 | 降低权重和激活的精度，从而降低内存要求。 | ✅ | ✅ | ❌ |
剪枝器 | 移除不重要或冗余的连接和神经元，从而形成更稀疏、更高效的网络。 | ✅ | ✅ | ❌ |
恢复器 | 恢复模型压缩后的性能。 | ➖ | ➖ | ✅ |
分解器 | 分解将几个小型矩阵乘法批量处理成一个大型融合操作。 | ✅ | ➖ | ➖ |
增强器 | 通过应用去噪或放大等后处理算法来增强模型输出。 | ❌ | - | ✅ |
✅ (提升), ➖ (大约相同), ❌ (恶化)
在 Pruna 文档中探索全部优化方法。
安装
使用以下命令安装 Pruna。
pip install pruna
优化 Diffusers 模型
Diffusers 模型支持多种优化算法，如下所示。
以下示例使用分解器、编译器和缓存器算法的组合来优化 black-forest-labs/FLUX.1-dev。这种组合将推理速度提高了 4.2 倍，并将峰值 GPU 内存使用量从 34.7GB 削减到 28.0GB，同时保持几乎相同的输出质量。
请参阅 Pruna 优化文档以了解此示例中使用的优化技术。
首先定义一个包含要使用的优化算法的 SmashConfig。要优化模型，请使用 smash 包装管道和 SmashConfig，然后像正常推理一样使用管道。
import torch
from diffusers import FluxPipeline
from pruna import PrunaModel, SmashConfig, smash
# load the model
# Try segmind/Segmind-Vega or black-forest-labs/FLUX.1-schnell with a small GPU memory
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
torch_dtype=torch.bfloat16
).to("cuda")
# define the configuration
smash_config = SmashConfig()
smash_config["factorizer"] = "qkv_diffusers"
smash_config["compiler"] = "torch_compile"
smash_config["torch_compile_target"] = "module_list"
smash_config["cacher"] = "fora"
smash_config["fora_interval"] = 2
# for the best results in terms of speed you can add these configs
# however they will increase your warmup time from 1.5 min to 10 min
# smash_config["torch_compile_mode"] = "max-autotune-no-cudagraphs"
# smash_config["quantizer"] = "torchao"
# smash_config["torchao_quant_type"] = "fp8dq"
# smash_config["torchao_excluded_modules"] = "norm+embedding"
# optimize the model
smashed_pipe = smash(pipe, smash_config)
# run the model
smashed_pipe("a knitted purple prune").images[0]
优化后，我们可以使用 Hugging Face Hub 共享和加载优化后的模型。
# save the model
smashed_pipe.save_to_hub("<username>/FLUX.1-dev-smashed")
# load the model
smashed_pipe = PrunaModel.from_hub("<username>/FLUX.1-dev-smashed")评估和基准测试 Diffusers 模型
Pruna 提供 EvaluationAgent 来评估优化模型的质量。
我们可以定义我们关心的指标，例如总时间和吞吐量，以及要评估的数据集。我们可以定义一个模型并将其传递给 EvaluationAgent。
优化模型
独立模型
我们可以使用 EvaluationAgent 加载和评估优化模型，并将其传递给 Task。
import torch
from diffusers import FluxPipeline
from pruna import PrunaModel
from pruna.data.pruna_datamodule import PrunaDataModule
from pruna.evaluation.evaluation_agent import EvaluationAgent
from pruna.evaluation.metrics import (
ThroughputMetric,
TorchMetricWrapper,
TotalTimeMetric,
)
from pruna.evaluation.task import Task
# define the device
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
# load the model
# Try PrunaAI/Segmind-Vega-smashed or PrunaAI/FLUX.1-dev-smashed with a small GPU memory
smashed_pipe = PrunaModel.from_hub("PrunaAI/FLUX.1-dev-smashed")
# Define the metrics
metrics = [
TotalTimeMetric(n_iterations=20, n_warmup_iterations=5),
ThroughputMetric(n_iterations=20, n_warmup_iterations=5),
TorchMetricWrapper("clip"),
]
# Define the datamodule
datamodule = PrunaDataModule.from_string("LAION256")
datamodule.limit_datasets(10)
# Define the task and evaluation agent
task = Task(metrics, datamodule=datamodule, device=device)
eval_agent = EvaluationAgent(task)
# Evaluate smashed model and offload it to CPU
smashed_pipe.move_to_device(device)
smashed_pipe_results = eval_agent.evaluate(smashed_pipe)
smashed_pipe.move_to_device("cpu")现在您已经了解了如何优化和评估模型，您可以开始使用 Pruna 优化您自己的模型。幸运的是，我们有许多示例可以帮助您入门。
有关 Flux 基准测试的更多详细信息，请查看宣布 FLUX-Juiced：最快的图像生成端点（快 2.6 倍）！博客文章和 InferBench Space。