使用🤗 Transformers优化Bark

用🤗 Transformers优化Bark

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

🤗 Hugging Face生态系统提供了一系列现成且易于使用的优化工具,可以应用于库中的所有模型。通过添加几行额外的代码,可以轻松地减少内存占用提高推理性能

在这个实践教程中,我将演示如何通过三个简单的优化来优化Bark,这是一个由🤗 Transformers支持的文本到语音(TTS)模型。这些优化仅依赖于🤗生态系统中的Transformers、Optimum和Accelerate库。

这个教程也是演示如何对一个非优化模型及其不同的优化进行基准测试的示例。

如果想要更简洁的教程版本,其中包含较少的解释但所有的代码,请参阅附带的Google Colab。

本文按照以下方式组织:

目录

  1. 关于Bark架构的提醒
  2. 不同优化技术及其优点的概述
  3. 基准测试结果的介绍

Bark是由Suno AI提出的基于transformer的文本到语音模型,可以生成各种音频输出,包括语音、音乐、背景噪音和简单的音效。此外,它还可以产生笑声、叹息和哭泣等非语言交流声音。

从v4.31.0版本开始,Bark已经在🤗 Transformers中可用!

您可以在这里尝试Bark并发现它的功能。

Bark由4个主要模型组成:

  • BarkSemanticModel(也称为“文本”模型):一个因果自回归变压器模型,以标记化文本作为输入,预测捕捉文本含义的语义文本标记。
  • BarkCoarseModel(也称为“粗略音韵学”模型):一个因果自回归变压器,以BarkSemanticModel模型的结果作为输入。它旨在预测EnCodec所需的前两个音频码本。
  • BarkFineModel(“精细音韵学”模型):这次是一个非因果自编码器变压器,它通过前面码本嵌入的总和来迭代地预测最后的码本。
  • 在预测了EncodecModel中的所有码本通道后,Bark使用它来解码输出音频数组。

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

加载模型和处理器

可以从Hugging Face Hub上的预训练权重加载预训练的Bark小和大检查点。您可以使用您希望使用的检查点大小更改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 = "让我们尝试使用Bark进行语音生成,它是一个文本到语音的模型"
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):

  # 定义测量生成过程开始和结束的事件
  start_event = torch.cuda.Event(enable_timing=True)
  end_event = torch.cuda.Event(enable_timing=True)

  # 重置CUDA内存统计并清空缓存
  torch.cuda.reset_peak_memory_stats(device)
  torch.cuda.empty_cache()
  torch.cuda.synchronize()

  # 获取开始时间
  start_event.record()

  # 实际生成
  for _ in range(nb_loops):
        # 设置种子以保证可重复性
        set_seed(0)
        output = model.generate(**inputs, do_sample = True, fine_temperature = 0.4, coarse_temperature = 0.8)

  # 获取结束时间
  end_event.record()
  torch.cuda.synchronize()

  # 测量内存占用和经过的时间
  max_memory = torch.cuda.max_memory_allocated(device)
  elapsed_time = start_event.elapsed_time(end_event) * 1.0e-3

  print('执行时间:', elapsed_time/nb_loops, '秒')
  print('最大内存占用:', max_memory*1e-9, ' GB')

  return output

基准情况

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

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

输出:

执行时间:9.3841625秒
最大内存占用:1.914612224 GB

现在,听听生成的语音输出:

from IPython.display import Audio

# 现在,听听输出
sampling_rate = model.generation_config.sample_rate
Audio(speech_output[0].cpu().numpy(), rate=sampling_rate)

输出听起来像这样(下载音频):

您的浏览器不支持 audio 元素。

重要说明:

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

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

这导致measure_latency_and_memory_use测量的延迟实际上可能不反映优化技术的实际性能!本文末尾的基准测试报告了在100次迭代中平均的结果,这是对模型性能的真实指示。

1. 🤗 更好的Transformer

更好的Transformer是🤗 Optimum 功能,它在内部执行内核融合。这意味着某些模型操作在GPU上的优化效果更好,模型的速度也更快。

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

通过一种名为Flash Attention的技术,可以对朴素的注意力机制进行大幅优化,该技术由Dao等人在2022年提出。

Flash Attention是一种更快、更高效的注意力计算算法,它结合了传统方法(如切片和重新计算)来最小化内存使用并提高速度。与以前的算法不同,Flash Attention将内存使用从二次降低为线性,这在对内存效率至关重要的应用中特别有用。

事实证明,🤗 更好的Transformer直接支持Flash Attention!只需要一行代码将模型导出为🤗 更好的Transformer并启用Flash Attention:

model = model.to_bettertransformer()

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

输出:

执行时间:5.43284375秒
最大内存占用:1.9151841280000002 GB

输出听起来像这样(下载音频):

您的浏览器不支持音频元素。

它能给我们带来什么?

性能没有降低,这意味着您可以获得与不使用此函数完全相同的结果,同时提高20%至30%的速度!想了解更多?请查看此博文。

2. 半精度

大多数AI模型通常使用一种称为单精度浮点数的存储格式,即 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)

输出:

执行时间:7.00045390625秒
最大内存占用:2.7436124160000004 GB

输出听起来像这样(下载音频):

您的浏览器不支持音频元素。

它能给我们带来什么?

在性能略微下降的情况下,您可以减少50%的内存占用,并获得5%的速度提升。

3. CPU卸载

正如本手册的第一部分所述,Bark包括4个子模型,这些子模型在音频生成期间被顺序调用。 换句话说,当一个子模型在使用时,其他子模型处于空闲状态。

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

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

好消息:在Bark中集成了CPU卸载功能,并且您只需使用一行代码即可使用它。

您只需要确保已安装🤗 Accelerate!

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

# 启用CPU卸载
model.enable_cpu_offload()

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

输出:

执行时间:8.97633828125秒
最大内存占用:1.3231160320000002 GB

输出听起来像这样(下载音频):

您的浏览器不支持音频元素。

它能给我们带来什么?

在速度略微下降(10%)的情况下,您可以获得巨大的内存占用减少(60% 🤯)。

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

想要更多吗?启用 fp16 后,内存占用甚至降至1GB。在下一部分中,我们将在实践中看到这一点!

4. 结合

让我们把所有东西都集合起来。好消息是,您可以结合优化技术,这意味着您可以使用CPU卸载,以及半精度和🤗 Better Transformer!

# 以fp16加载模型
model = BarkModel.from_pretrained("suno/bark-small", torch_dtype=torch.float16).to(device)

# 转换为BetterTransformer模型
model = BetterTransformer.transform(model, keep_original_model=False)

# 启用CPU offload
model.enable_cpu_offload()

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

输出:

执行时间: 7.4496484375000005 秒
最大内存占用:0.46871091200000004 GB

输出的声音如下(下载音频):

您的浏览器不支持音频元素。

这带来了什么好处?

最终,您可以获得23%的加速和80%的内存节省!

使用批处理

还想要更多吗?

总的来说,当进行批处理时,这3种优化技术带来的结果更好。批处理意味着将多个样本的操作组合在一起,以使生成样本的总时间低于逐个生成样本。

这里是一个快速示例,展示如何使用批处理:

text_prompt = [
    "让我们尝试使用Bark生成语音,这是一个文本转语音模型",
    "哇,批处理太棒了!",
    "我喜欢Hugging Face,它太酷了。"]

inputs = processor(text_prompt).to(device)

with torch.inference_mode():
  # 一次性生成所有样本
  speech_output = model.generate(**inputs, do_sample = True, fine_temperature = 0.4, coarse_temperature = 0.8)

输出的声音如下(下载第一个、第二个和最后一个音频):

您的浏览器不支持音频元素。您的浏览器不支持音频元素。您的浏览器不支持音频元素。

如上所述,我们进行的小实验只是一种思考的练习,并需要进行扩展以更好地衡量性能。在正确测量性能之前,还需要在几个空白迭代中对GPU进行预热。

下面是使用Bark的大版本进行100个样本基准测试的结果。

该基准测试在一台具有256个新令牌的NVIDIA TITAN RTX 24GB上运行。

如何读取结果?

延迟

它测量生成方法的单个调用的持续时间,无论批处理大小如何。

换句话说,它等于elapsedTimenbLoops\frac{elapsedTime}{nbLoops}nbLoopselapsedTime​。

更低的延迟更好。

最大内存占用

它测量在单个调用生成方法期间使用的最大内存。

更低的内存占用更好。

吞吐量

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

换句话说,它等于nbLoops∗batchSizeelapsedTime\frac{nbLoops*batchSize}{elapsedTime}elapsedTimenbLoops∗batchSize​。

更高的吞吐量更好。

无批处理

以下是batch_size=1的结果。

评论

如预期,CPU offload大大减少了内存占用,同时稍微增加了延迟。

然而,结合了bettertransformer和fp16,我们同时获得了延迟和内存的巨大减少!

批处理大小设置为8

这里是batch_size=8的基准测试结果和吞吐量测量。

请注意,由于bettertransformer是一个免费的优化,因为它执行的操作与非优化模型完全相同,内存占用也相同,但速度更快,所以基准测试默认启用了该优化。

评论

这里我们可以看到将这三个优化特性结合起来的潜力!

fp16对延迟的影响在batch_size = 1时不太明显,但在这里却非常有趣,因为它可以减少近一半的延迟,并且几乎可以将吞吐量翻倍!

本博文展示了几个简单的优化技巧,它们打包在🤗生态系统中。使用这些技术中的任何一种,或者将它们三者结合起来使用,都可以极大地提高Bark推断速度和内存占用。

  • 你可以使用Bark的大版本,而不会有任何性能下降,内存占用仅为2GB,而不是5GB,速度比原来提高了15%,使用🤗更好的Transformer和CPU卸载

  • 你更喜欢高吞吐量吗?使用🤗更好的Transformer和半精度,每批次8个

  • 你可以同时获得两全其美的效果,使用fp16,🤗更好的Transformer和CPU卸载