Diffusers 文件
Pruna
並獲得增強的文件體驗
開始使用
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。