使用 🤗 Transformers 优化文本转语音模型

🤗 Transformers 提供了许多领域和任务中最先进（SoTA）的模型。为了从这些模型中获得最佳性能，需要对它们进行优化，以提高推理速度和降低内存使用量。

🤗 Hugging Face 生态系统提供了现成易用的优化工具，可以全面应用于库中的所有模型。这使得仅需少量额外代码即可轻松**减少内存占用**并**改善推理**。

在本实战教程中，我将演示如何基于三种简单优化来优化 🤗 Transformers 支持的文本转语音（TTS）模型 Bark。这些优化仅依赖于 🤗 生态系统中的 Transformers、Optimum 和 Accelerate 库。

本教程也演示了如何对未优化模型及其各种优化进行基准测试。

如需更精简的版本（解释较少但包含所有代码），请参阅随附的 Google Colab。

本博文的组织结构如下：

目录

1. Bark 架构 概述
2. 不同优化技术及其优点的 概述
3. 基准测试结果的 展示

Bark 架构

Bark 是 Suno AI 在 suno-ai/bark 中提出的基于 Transformer 的文本转语音模型。它能够生成各种音频输出，包括语音、音乐、背景噪音和简单的音效。此外，它还能生成非语言交流声音，如笑声、叹息声和哭泣声。

Bark 自 v4.31.0 版本起已可在 🤗 Transformers 中使用！

您可以在这里试用 Bark 并探索其功能。

Bark 由 4 个主要模型组成

• BarkSemanticModel（也称为“文本”模型）：一个因果自回归 Transformer 模型，将分词后的文本作为输入，并预测捕获文本含义的语义文本 token。
• BarkCoarseModel（也称为“粗声学”模型）：一个因果自回归 Transformer，以 BarkSemanticModel 模型的输出作为输入。它旨在预测 EnCodec 所需的前两个音频码本。
• BarkFineModel（“细声学”模型），这次是一个非因果自编码器 Transformer，它根据先前码本嵌入的总和迭代地预测最后一个码本。
• 在预测了 EncodecModel 的所有码本通道后，Bark 使用它来解码输出音频数组。

撰写本文时，有两个 Bark 检查点可用，一个较小版本和一个较大版本。

加载模型及其处理器

预训练的 Bark 小版本和大版本检查点可以从 Hugging Face Hub 上的预训练权重加载。您可以根据需要更改 repo-id 来选择检查点大小。

我们将默认使用小型检查点，以保持快速。但您可以通过使用 ` "suno/bark"` 而不是 ` "suno/bark-small"` 来尝试大型检查点。

from transformers import BarkModel

model = BarkModel.from_pretrained("suno/bark-small")

将模型放置到加速器设备以充分利用优化技术

import torch

device = "cuda:0" if torch.cuda.is_available() else "cpu"
model = model.to(device)

加载处理器，它将负责分词和可选的说话人嵌入。

from transformers import AutoProcessor

processor = AutoProcessor.from_pretrained("suno/bark-small")

优化技术

在本节中，我们将探讨如何使用 🤗 Optimum 和 🤗 Accelerate 库中现成的功能，以最少的代码更改来优化 Bark 模型。

一些设置

让我们准备输入并定义一个函数来测量 Bark 生成方法的延迟和 GPU 内存占用。

text_prompt = "Let's try generating speech, with Bark, a text-to-speech model"
inputs = processor(text_prompt).to(device)

测量延迟和 GPU 内存占用需要使用特定的 CUDA 方法。我们定义了一个实用函数，用于测量模型在推理时的延迟和 GPU 内存占用。为了确保我们准确了解这些指标，我们对指定数量的运行 `nb_loops` 进行平均。

import torch
from transformers import set_seed


def measure_latency_and_memory_use(model, inputs, nb_loops = 5):

  # define Events that measure start and end of the generate pass
  start_event = torch.cuda.Event(enable_timing=True)
  end_event = torch.cuda.Event(enable_timing=True)

  # reset cuda memory stats and empty cache
  torch.cuda.reset_peak_memory_stats(device)
  torch.cuda.empty_cache()
  torch.cuda.synchronize()

  # get the start time
  start_event.record()

  # actually generate
  for _ in range(nb_loops):
        # set seed for reproducibility
        set_seed(0)
        output = model.generate(**inputs, do_sample = True, fine_temperature = 0.4, coarse_temperature = 0.8)

  # get the end time
  end_event.record()
  torch.cuda.synchronize()

  # measure memory footprint and elapsed time
  max_memory = torch.cuda.max_memory_allocated(device)
  elapsed_time = start_event.elapsed_time(end_event) * 1.0e-3

  print('Execution time:', elapsed_time/nb_loops, 'seconds')
  print('Max memory footprint', max_memory*1e-9, ' GB')

  return output

基本情况

在进行任何优化之前，让我们测量基线模型的性能并听一个生成的示例。我们将对模型进行五次迭代的基准测试，并报告这些指标的平均值。

with torch.inference_mode():
  speech_output = measure_latency_and_memory_use(model, inputs, nb_loops = 5)

输出

Execution time: 9.3841625 seconds
Max memory footprint 1.914612224  GB

现在，听一下输出。

from IPython.display import Audio

# now, listen to the output
sampling_rate = model.generation_config.sample_rate
Audio(speech_output[0].cpu().numpy(), rate=sampling_rate)

输出声音如下（下载音频）

重要提示：

这里，迭代次数实际上相当低。为了准确测量和比较结果，应该将其增加到至少 100 次。

增加 `nb_loops` 的主要原因之一是，即使输入固定，生成的语音长度在不同迭代之间也会有很大差异。

这样做的结果之一是，`measure_latency_and_memory_use` 测量的延迟可能无法真实反映优化技术的实际性能！博客文章末尾的基准测试报告了 100 次迭代的平均结果，这真实反映了模型的性能。

1. 🤗 Better Transformer

Better Transformer 是 🤗 Optimum 的一个功能，它在底层执行内核融合。这意味着某些模型操作将在 GPU 上得到更好的优化，并且模型最终会更快。

更具体地说，大多数受 🤗 Transformers 支持的模型都依赖于注意力机制，这使得它们在生成输出时能够有选择地关注输入的某些部分。这使得模型能够有效地处理长距离依赖关系并捕获数据中复杂的上下文关系。

原始的注意力技术可以通过 Dao 等人于 2022 年提出的 Flash Attention 技术得到极大优化。

Flash Attention 是一种更快、更高效的注意力计算算法，它结合了传统方法（如平铺和重新计算）以最大限度地减少内存使用并提高速度。与以前的算法不同，Flash Attention 将内存使用从序列长度的二次方降低到线性，这对于内存效率重要的应用特别有用。

事实证明，Flash Attention 被 🤗 Better Transformer 开箱即用支持！它只需要一行代码就可以将模型导出到 🤗 Better Transformer 并启用 Flash Attention

model =  model.to_bettertransformer()

with torch.inference_mode():
  speech_output = measure_latency_and_memory_use(model, inputs, nb_loops = 5)

输出

Execution time: 5.43284375 seconds
Max memory footprint 1.9151841280000002  GB

输出声音如下（下载音频）

它带来了什么？

性能没有下降，这意味着您可以在不使用此函数的情况下获得完全相同的结果，同时速度提高 20% 到 30%！想了解更多？请参阅这篇博客文章。

2. 半精度

大多数人工智能模型通常使用一种称为单精度浮点数的存储格式，即 fp32。这在实践中意味着什么？每个数字都使用 32 位存储。

因此，您可以选择使用 16 位对数字进行编码，这被称为半精度浮点数，即 fp16，并且使用的存储空间是以前的一半！不仅如此，您还可以获得推理速度的提升！

当然，这也会带来轻微的性能下降，因为模型内部的操作不会像使用 fp32 那样精确。

您可以通过在 `BarkModel.from_pretrained(...)` 行中简单地添加 `torch_dtype=torch.float16` 来加载半精度的 🤗 Transformers 模型！

换句话说，

model = BarkModel.from_pretrained("suno/bark-small", torch_dtype=torch.float16).to(device)

with torch.inference_mode():
  speech_output = measure_latency_and_memory_use(model, inputs, nb_loops = 5)

输出

Execution time: 7.00045390625 seconds
Max memory footprint 2.7436124160000004  GB

输出声音如下（下载音频）

它带来了什么？

性能略有下降，但内存占用减少了 50%，速度提升了 5%。

3. CPU 卸载

如本手册第一部分所述，Bark 由 4 个子模型组成，它们在音频生成过程中按顺序调用。**换句话说，当一个子模型在使用时，其他子模型处于空闲状态。**

这为什么是个问题？AI 中 GPU 内存非常宝贵，因为它是在其中操作最快的地方，而且它通常是一个瓶颈。

一个简单的解决方案是在不活动时将子模型从 GPU 卸载。此操作称为 CPU 卸载。

好消息：Bark 的 CPU 卸载已集成到 🤗 Transformers 中，您只需一行代码即可使用它。

您只需要确保安装了 🤗 Accelerate！

model = BarkModel.from_pretrained("suno/bark-small")

# Enable CPU offload
model.enable_cpu_offload()

with torch.inference_mode():
  speech_output = measure_latency_and_memory_use(model, inputs, nb_loops = 5)

输出

Execution time: 8.97633828125 seconds
Max memory footprint 1.3231160320000002  GB

输出声音如下（下载音频）

它带来了什么？

速度略有下降（10%），但内存占用大幅减少（60% 🤯）。

启用此功能后，bark-large 的占用空间现在仅为 2GB，而不是 5GB。这与 bark-small 的内存占用相同！

想了解更多？启用 fp16 后，甚至可以降至 1GB。我们将在下一节中看到它的实际应用！

4. 结合

让我们把所有这些结合起来。好消息是您可以组合优化技术，这意味着您可以使用 CPU 卸载、半精度和 🤗 Better Transformer！

# load in fp16
model = BarkModel.from_pretrained("suno/bark-small", torch_dtype=torch.float16).to(device)

# convert to bettertransformer
model = BetterTransformer.transform(model, keep_original_model=False)

# enable CPU offload
model.enable_cpu_offload()

with torch.inference_mode():
  speech_output = measure_latency_and_memory_use(model, inputs, nb_loops = 5)

输出

Execution time: 7.4496484375000005 seconds
Max memory footprint 0.46871091200000004  GB

输出声音如下（下载音频）

它带来了什么？

最终，您将获得 23% 的速度提升和高达 80% 的内存节省！

使用批处理

想要更多？

总而言之，当批处理时，这三种优化技术带来了更好的结果。批处理意味着将多个样本的操作组合起来，以使生成样本的总时间低于逐个生成样本。

这里有一个如何使用它的快速示例

text_prompt = [
    "Let's try generating speech, with Bark, a text-to-speech model",
    "Wow, batching is so great!",
    "I love Hugging Face, it's so cool."]

inputs = processor(text_prompt).to(device)


with torch.inference_mode():
  # samples are generated all at once
  speech_output = model.generate(**inputs, do_sample = True, fine_temperature = 0.4, coarse_temperature = 0.8)

输出声音如下（下载第一段、第二段和最后一段音频）

基准测试结果

如上所述，我们进行的这个小实验是一个思维练习，需要扩展以更好地衡量性能。在正确测量性能之前，还需要用几次空迭代来预热 GPU。

以下是针对 100 个样本的基准测试结果，使用 Bark 的大版本。

基准测试在 NVIDIA TITAN RTX 24GB 上运行，最大新 token 数为 256。

如何解读结果？

延迟

它衡量单次调用生成方法所需的时间，与批处理大小无关。

换句话说，它等于 $\frac{elapsedTime}{nbLoops}$。

延迟越低越好。

最大内存占用

它衡量单次调用生成方法期间使用的最大内存。

占用越低越好。

吞吐量

它衡量每秒生成的样本数量。这次考虑了批处理大小。

换句话说，它等于 $\frac{nbLoops*batchSize}{elapsedTime}$。

吞吐量越高越好。

不批处理

以下是 `batch_size=1` 的结果。

评论

正如所料，CPU 卸载大大减少了内存占用，同时略微增加了延迟。

然而，结合 Better Transformer 和 fp16，我们获得了两全其美的效果，大幅减少了延迟和内存！

批处理大小设置为 8

以下是基准测试结果，但使用了 `batch_size=8` 和吞吐量测量。

请注意，由于 `bettertransformer` 是一种免费的优化，因为它执行完全相同的操作，并且与未优化模型具有相同的内存占用，同时速度更快，因此基准测试默认**启用了此优化**。

评论

这就是我们看到结合所有三种优化功能的潜力的地方！

fp16 对延迟的影响在 batch_size = 1 时不太明显，但在这里它具有巨大的意义，因为它可以将延迟减少近一半，并将吞吐量几乎翻倍！

总结

这篇博文展示了 🤗 生态系统中捆绑的一些简单优化技巧。使用这些技术中的任何一种，或三者结合，都可以大大提高 Bark 的推理速度和内存占用。

• 您可以使用 Bark 的大型版本，而不会有任何性能下降，内存占用仅为 2GB（而不是 5GB），速度提高 15%，**使用 🤗 Better Transformer 和 CPU 卸载**。

• 您更喜欢高吞吐量吗？**使用 🤗 Better Transformer 和半精度进行 8 次批处理**。

• 通过使用 **fp16、🤗 Better Transformer 和 CPU 卸载**，您可以获得两全其美的效果！