使用 🤗 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 解除安裝**，您可以獲得兩全其美的效果！