稳定扩散推理加速技巧分析

核心观点与关键数据

稳定扩散与 TensorRT 结合

• 目标：利用 TensorRT 加速稳定扩散模型，特别是针对安培架构的 2:4 结构化稀疏模型。
• 方法：通过 PyTorch Sparsity API 生成 2:4 结构化稀疏模型，支持稀疏梯度反向传播，用于研究/微调。

稀疏矩阵压缩格式

• 压缩方式：在 4 个值的块中不存储两个 0，如果一个块包含两个以上的 0，部分 0 不存储。
• 元数据：使用索引元数据访问剩余的 2 个值，每个值 2 位，fp16 开销 12.5%，int8 开销 25%。
• 工具：NVIDIA APEX 库（包括 AMP 和 ASP）支持混合精度和稀疏训练。

TensorRT 部署

• 流程：PyTorch 检查点 -> ONNX -> TensorRT 引擎。
• 命令：trtexec --onnx=model.onnx --saveEngine=engine.trt --explicitBatch --sparsity=enable --fp16。

稀疏基准测试

• 模型：SD 1.5 和 SD 2.1，宽度为 512x512。
• 结果：展示了稀疏性训练后的验证图像，表明模型在稀疏化后仍保持较高性能。

量化与重新安装

• 量化目标：模型架构量化。
• 步骤：LoRA TensorRT 重新安装，对齐 ONNX 权重名称与 PyTorch 权重。