使用 🤗 Transformers 优化 Bark

Hugging Face 已详细介绍了一种优化 Bark 文本转语音（TTS）模型的方法，以降低内存使用并提高推理速度。通过结合三种特定的优化技术——Better Transformer、半精度浮点数（fp16）和 CPU 卸载——用户可以实现高达 80% 的内存占用减少和 23% 的速度提升。

Bark 模型架构

Bark 是由 Suno AI 开发的基于 transformer 的 TTS 模型，能够生成语音、音乐、背景噪声和非语言声音（例如笑声和叹息）。其架构由四个按顺序运行的主要组件组成：

  • BarkSemanticModel：一个因果自回归 transformer，用于预测语义文本标记。
  • BarkCoarseModel：一个因果自回归 transformer，用于预测 EnCodec 的前两个音频码本。
  • BarkFineModel：一个非因果自编码 transformer，迭代预测剩余的码本。
  • EncodecModel：将最终码本通道解码为输出音频数组。

优化技术

通过 🤗 Optimum 和 🤗 Accelerate 库实现优化，仅需最少的代码更改。

Better Transformer

Better Transformer 利用内核融合和 Flash Attention 来优化 GPU 操作。Flash Attention 将内存使用从序列长度的二次关系降低到线性关系，从而在不降低模型性能的情况下提高速度。

主要优势：在不损失输出质量的情况下提供 20% 至 30% 的速度提升。

半精度（fp16）

通过将单精度浮点数（fp32，每个数字 32 位）切换为半精度（fp16，每个数字 16 位），模型的存储需求减半。

主要优势：将内存占用减少 50%，并提供适度的速度提升（大约 5%），尽管可能会引入轻微的性能下降。

CPU 卸载

由于 Bark 的四个子模型是按顺序调用的，同一时间只有一个处于活动状态，而其他保持空闲。CPU 卸载将不活动的子模型从 GPU 卸载到 CPU，以释放宝贵的 GPU 内存。

主要优势：将内存占用减少约 60%，同时延迟略增（约 10%）。

性能基准

基准测试在 NVIDIA TITAN RTX 24GB 上使用 Bark 的大型版本进行，最大生成 256 个新 token，平均取 100 个样本的结果。

单样本生成（批量大小 = 1）

生成单个样本时，Better Transformer、CPU 卸载和 fp16 的组合带来最显著的收益：

优化 延迟变化 内存变化
未优化 0% 0%
仅 Better Transformer -27% -1%
卸载 + Better Transformer -15% -59%
卸载 + Better Transformer + fp16 -23% -80%

批量生成（批量大小 = 8）

将多个样本组合在一起可显著提高吞吐量。在默认启用 Better Transformer 的情况下，观察到以下结果：

优化 延迟变化 内存变化 吞吐量变化
基准情况（Better Transformer） 0% 0% 0%
+ fp16 -46% -50% +87%
+ 卸载 +6% -38% -6%
+ 卸载 + fp16 -43% -69% +77%

最优配置摘要

根据使用场景，推荐不同的优化组合：

  • 为了内存效率：使用 Better Transformer 和 CPU 卸载，以 2GB 的占用运行大型 Bark 模型，而不是 5GB。
  • 为了高吞吐量：使用批量（大小 8）结合 Better Transformer 和半精度（fp16）。
  • 为了平衡性能：结合 fp16、Better Transformer 和 CPU 卸载，以在延迟和内存使用方面获得最佳整体降低。

Sources