注意： TensorFlow Lite 现在是 Google AI Edge 的一部分。最新文档现已在 ai.google.dev/edge/lite 上提供。了解更多

使用 TensorFlow Lite Model Maker 重新训练语音识别模型

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

在本 Colab 笔记本中，您将学习如何使用 TensorFlow Lite Model Maker 训练一个语音识别模型，该模型可以使用一秒钟的音频样本对口语单词或短语进行分类。Model Maker 库使用迁移学习来使用新数据集重新训练现有的 TensorFlow 模型，这减少了训练所需的样本数据量和时间。

默认情况下，此笔记本使用 TFJS 语音命令识别器中的单词子集（例如“向上”、“向下”、“向左”和“向右”）重新训练模型（BrowserFft）。然后，它导出一个 TFLite 模型，您可以在移动设备或嵌入式系统（例如 Raspberry Pi）上运行它。它还将训练后的模型导出为 TensorFlow SavedModel。

此笔记本还设计为接受自定义的 WAV 文件数据集，这些文件以 ZIP 文件的形式上传到 Colab。每个类别拥有的样本越多，准确率就越高，但由于迁移学习过程使用来自预训练模型的特征嵌入，因此您仍然可以使用每个类别中只有几十个样本获得相当准确的模型。

如果您想使用默认语音数据集运行笔记本，现在就可以通过单击 Colab 工具栏中的运行时 > 全部运行来运行整个笔记本。但是，如果您想使用自己的数据集，请继续向下滚动到准备数据集并按照那里的说明操作。

导入所需的包

您将需要 TensorFlow、TFLite Model Maker 以及一些用于音频操作、播放和可视化的模块。

sudo apt -y install libportaudio2
pip install tflite-model-maker

import os
import glob
import random
import shutil

import librosa
import soundfile as sf
from IPython.display import Audio
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

import tensorflow as tf
import tflite_model_maker as mm
from tflite_model_maker import audio_classifier
from tflite_model_maker.config import ExportFormat

print(f"TensorFlow Version: {tf.__version__}")
print(f"Model Maker Version: {mm.__version__}")

准备数据集

要使用默认语音数据集进行训练，只需按原样运行以下所有代码。

但如果您想使用自己的语音数据集进行训练，请按照以下步骤操作

确保您的数据集中每个样本都以 **WAV 文件格式，时长约为一秒钟**。然后创建一个包含所有 WAV 文件的 ZIP 文件，并将它们组织到每个分类的单独子文件夹中。例如，每个语音命令“yes”的样本都应放在名为“yes”的子文件夹中。即使您只有一个类别，样本也必须保存在以类别名称命名的子目录中。（此脚本假设您的数据集 **未拆分** 为训练/验证/测试集，并为您执行此拆分。）
单击左侧面板中的 **文件** 选项卡，然后将您的 ZIP 文件拖放到那里以上传。
使用以下下拉选项将 **use_custom_dataset** 设置为 True。
然后跳到准备自定义音频数据集以指定您的 ZIP 文件名和数据集目录名。

use_custom_dataset = False

生成背景噪音数据集

无论您使用的是默认语音数据集还是自定义数据集，您都应该有一组良好的背景噪音，以便您的模型能够区分语音和其他噪音（包括静音）。

由于以下背景样本以 WAV 文件形式提供，时长为一分钟或更长，因此我们需要将它们拆分为更小的、一秒钟的样本，以便我们可以为测试数据集保留一些样本。我们还将结合几个不同的样本来源来构建一组全面的背景噪音和静音

tf.keras.utils.get_file('speech_commands_v0.01.tar.gz',
                        'http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz',
                        cache_dir='./',
                        cache_subdir='dataset-speech',
                        extract=True)
tf.keras.utils.get_file('background_audio.zip',
                        'https://storage.googleapis.com/download.tensorflow.org/models/tflite/sound_classification/background_audio.zip',
                        cache_dir='./',
                        cache_subdir='dataset-background',
                        extract=True)

# Create a list of all the background wav files
files = glob.glob(os.path.join('./dataset-speech/_background_noise_', '*.wav'))
files = files + glob.glob(os.path.join('./dataset-background', '*.wav'))

background_dir = './background'
os.makedirs(background_dir, exist_ok=True)

# Loop through all files and split each into several one-second wav files
for file in files:
  filename = os.path.basename(os.path.normpath(file))
  print('Splitting', filename)
  name = os.path.splitext(filename)[0]
  rate = librosa.get_samplerate(file)
  length = round(librosa.get_duration(filename=file))
  for i in range(length - 1):
    start = i * rate
    stop = (i * rate) + rate
    data, _ = sf.read(file, start=start, stop=stop)
    sf.write(os.path.join(background_dir, name + str(i) + '.wav'), data, rate)

准备语音命令数据集

我们已经下载了语音命令数据集，所以现在我们只需要修剪模型的类别数量。

此数据集包含 30 多个语音命令分类，其中大多数分类有 2,000 多个样本。但是，由于我们使用的是迁移学习，因此我们不需要那么多样本。所以以下代码做了一些事情

指定我们要使用的分类，并删除其余分类。
仅保留每个类别 150 个样本用于训练（以证明迁移学习在较小的数据集上也能很好地工作，并且只是为了减少训练时间）。
为测试数据集创建一个单独的目录，以便我们以后可以轻松地用它们运行推理。

if not use_custom_dataset:
  commands = [ "up", "down", "left", "right", "go", "stop", "on", "off", "background"]
  dataset_dir = './dataset-speech'
  test_dir = './dataset-test'

  # Move the processed background samples
  shutil.move(background_dir, os.path.join(dataset_dir, 'background'))   

  # Delete all directories that are not in our commands list
  dirs = glob.glob(os.path.join(dataset_dir, '*/'))
  for dir in dirs:
    name = os.path.basename(os.path.normpath(dir))
    if name not in commands:
      shutil.rmtree(dir)

  # Count is per class
  sample_count = 150
  test_data_ratio = 0.2
  test_count = round(sample_count * test_data_ratio)

  # Loop through child directories (each class of wav files)
  dirs = glob.glob(os.path.join(dataset_dir, '*/'))
  for dir in dirs:
    files = glob.glob(os.path.join(dir, '*.wav'))
    random.seed(42)
    random.shuffle(files)
    # Move test samples:
    for file in files[sample_count:sample_count + test_count]:
      class_dir = os.path.basename(os.path.normpath(dir))
      os.makedirs(os.path.join(test_dir, class_dir), exist_ok=True)
      os.rename(file, os.path.join(test_dir, class_dir, os.path.basename(file)))
    # Delete remaining samples
    for file in files[sample_count + test_count:]:
      os.remove(file)

准备自定义数据集

如果您想用我们自己的语音数据集训练模型，您需要将您的样本作为 WAV 文件上传到 ZIP 文件中 (如上所述)，并修改以下变量以指定您的数据集

if use_custom_dataset:
  # Specify the ZIP file you uploaded:
  !unzip YOUR-FILENAME.zip
  # Specify the unzipped path to your custom dataset
  # (this path contains all the subfolders with classification names):
  dataset_dir = './YOUR-DIRNAME'

更改完上面的文件名和路径名后，您就可以使用自定义数据集训练模型了。在 Colab 工具栏中，选择 **运行时 > 运行全部** 以运行整个笔记本。

以下代码将我们的新背景噪音样本集成到您的数据集中，然后分离一部分样本以创建测试集。

def move_background_dataset(dataset_dir):
  dest_dir = os.path.join(dataset_dir, 'background')
  if os.path.exists(dest_dir):
    files = glob.glob(os.path.join(background_dir, '*.wav'))
    for file in files:
      shutil.move(file, dest_dir)
  else:
    shutil.move(background_dir, dest_dir)

if use_custom_dataset:
  # Move background samples into custom dataset
  move_background_dataset(dataset_dir)

  # Now we separate some of the files that we'll use for testing:
  test_dir = './dataset-test'
  test_data_ratio = 0.2
  dirs = glob.glob(os.path.join(dataset_dir, '*/'))
  for dir in dirs:
    files = glob.glob(os.path.join(dir, '*.wav'))
    test_count = round(len(files) * test_data_ratio)
    random.seed(42)
    random.shuffle(files)
    # Move test samples:
    for file in files[:test_count]:
      class_dir = os.path.basename(os.path.normpath(dir))
      os.makedirs(os.path.join(test_dir, class_dir), exist_ok=True)
      os.rename(file, os.path.join(test_dir, class_dir, os.path.basename(file)))
    print('Moved', test_count, 'images from', class_dir)

播放样本

为了确保数据集看起来正确，让我们从测试集中播放一个随机样本

def get_random_audio_file(samples_dir):
  files = os.path.abspath(os.path.join(samples_dir, '*/*.wav'))
  files_list = glob.glob(files)
  random_audio_path = random.choice(files_list)
  return random_audio_path

def show_sample(audio_path):
  audio_data, sample_rate = sf.read(audio_path)
  class_name = os.path.basename(os.path.dirname(audio_path))
  print(f'Class: {class_name}')
  print(f'File: {audio_path}')
  print(f'Sample rate: {sample_rate}')
  print(f'Sample length: {len(audio_data)}')

  plt.title(class_name)
  plt.plot(audio_data)
  display(Audio(audio_data, rate=sample_rate))

random_audio = get_random_audio_file(test_dir)
show_sample(random_audio)

定义模型

使用 Model Maker 重新训练任何模型时，您必须首先定义一个模型规范。规范定义了基础模型，您的新模型将从中提取特征嵌入以开始学习新类别。此语音识别器的规范基于预训练的 TFJS 的 BrowserFft 模型。

该模型期望输入为一个音频样本，采样率为 44.1 kHz，时长略低于一秒钟：确切的样本长度必须为 44034 帧。

您不需要对训练数据集进行任何重采样。Model Maker 会为您处理这些操作。但是，当您稍后运行推理时，必须确保您的输入与预期格式匹配。

您只需要实例化 BrowserFftSpec

spec = audio_classifier.BrowserFftSpec()

加载您的数据集

现在您需要根据模型规范加载您的数据集。Model Maker 包含 DataLoader API，它将从文件夹中加载您的数据集，并确保它处于模型规范的预期格式。

我们已经通过将一些测试文件移动到单独的目录中来保留了一些测试文件，这使得以后用它们运行推理变得更加容易。现在我们将为每个拆分创建一个 DataLoader：训练集、验证集和测试集。

加载语音命令数据集

if not use_custom_dataset:
  train_data_ratio = 0.8
  train_data = audio_classifier.DataLoader.from_folder(
      spec, dataset_dir, cache=True)
  train_data, validation_data = train_data.split(train_data_ratio)
  test_data = audio_classifier.DataLoader.from_folder(
      spec, test_dir, cache=True)

加载自定义数据集

if use_custom_dataset:
  train_data_ratio = 0.8
  train_data = audio_classifier.DataLoader.from_folder(
      spec, dataset_dir, cache=True)
  train_data, validation_data = train_data.split(train_data_ratio)
  test_data = audio_classifier.DataLoader.from_folder(
      spec, test_dir, cache=True)

训练模型

现在我们将使用 Model Maker 的 create() 函数根据我们的模型规范和训练数据集创建模型，并开始训练。

如果您使用的是自定义数据集，您可能需要根据训练集中样本的数量相应地更改批次大小。

# If your dataset has fewer than 100 samples per class,
# you might want to try a smaller batch size
batch_size = 25
epochs = 25
model = audio_classifier.create(train_data, spec, validation_data, batch_size, epochs)

查看模型性能

即使从上面的训练输出中可以看出准确率/损失看起来很好，但重要的是也要使用模型尚未见过的测试数据运行模型，这就是 evaluate() 方法在这里所做的。

model.evaluate(test_data)

查看混淆矩阵

在训练像这样的分类模型时，检查混淆矩阵也很有用。混淆矩阵为您提供了分类器在测试数据中对每个分类的执行情况的详细可视化表示。

def show_confusion_matrix(confusion, test_labels):
  """Compute confusion matrix and normalize."""
  confusion_normalized = confusion.astype("float") / confusion.sum(axis=1)
  sns.set(rc = {'figure.figsize':(6,6)})
  sns.heatmap(
      confusion_normalized, xticklabels=test_labels, yticklabels=test_labels,
      cmap='Blues', annot=True, fmt='.2f', square=True, cbar=False)
  plt.title("Confusion matrix")
  plt.ylabel("True label")
  plt.xlabel("Predicted label")

confusion_matrix = model.confusion_matrix(test_data)
show_confusion_matrix(confusion_matrix.numpy(), test_data.index_to_label)

导出模型

最后一步是将您的模型导出为 TensorFlow Lite 格式，以便在移动/嵌入式设备上执行，并导出为 SavedModel 格式以便在其他地方执行。

从 Model Maker 导出 .tflite 文件时，它包含模型元数据，描述了以后在推理期间可能会有所帮助的各种详细信息。它甚至包含分类标签文件的副本，因此您不需要单独的 labels.txt 文件。（在下一节中，我们将展示如何使用此元数据来运行推理。）

TFLITE_FILENAME = 'browserfft-speech.tflite'
SAVE_PATH = './models'

print(f'Exporing the model to {SAVE_PATH}')
model.export(SAVE_PATH, tflite_filename=TFLITE_FILENAME)
model.export(SAVE_PATH, export_format=[mm.ExportFormat.SAVED_MODEL, mm.ExportFormat.LABEL])

使用 TF Lite 模型运行推理

现在，您的 TFLite 模型可以使用任何支持的推理库或使用新的 TFLite AudioClassifier 任务 API 部署和运行。以下代码展示了如何在 Python 中使用 .tflite 模型运行推理。

# This library provides the TFLite metadata API
 pip install -q tflite_support

from tflite_support import metadata
import json

def get_labels(model):
  """Returns a list of labels, extracted from the model metadata."""
  displayer = metadata.MetadataDisplayer.with_model_file(model)
  labels_file = displayer.get_packed_associated_file_list()[0]
  labels = displayer.get_associated_file_buffer(labels_file).decode()
  return [line for line in labels.split('\n')]

def get_input_sample_rate(model):
  """Returns the model's expected sample rate, from the model metadata."""
  displayer = metadata.MetadataDisplayer.with_model_file(model)
  metadata_json = json.loads(displayer.get_metadata_json())
  input_tensor_metadata = metadata_json['subgraph_metadata'][0][
          'input_tensor_metadata'][0]
  input_content_props = input_tensor_metadata['content']['content_properties']
  return input_content_props['sample_rate']

要观察模型在真实样本上的表现如何，请反复运行以下代码块。每次运行时，它都会获取一个新的测试样本并用它运行推理，您可以在下面听到音频样本。

# Get a WAV file for inference and list of labels from the model
tflite_file = os.path.join(SAVE_PATH, TFLITE_FILENAME)
labels = get_labels(tflite_file)
random_audio = get_random_audio_file(test_dir)

# Ensure the audio sample fits the model input
interpreter = tf.lite.Interpreter(tflite_file)
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
input_size = input_details[0]['shape'][1]
sample_rate = get_input_sample_rate(tflite_file)
audio_data, _ = librosa.load(random_audio, sr=sample_rate)
if len(audio_data) < input_size:
  audio_data.resize(input_size)
audio_data = np.expand_dims(audio_data[:input_size], axis=0)

# Run inference
interpreter.allocate_tensors()
interpreter.set_tensor(input_details[0]['index'], audio_data)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])

# Display prediction and ground truth
top_index = np.argmax(output_data[0])
label = labels[top_index]
score = output_data[0][top_index]
print('---prediction---')
print(f'Class: {label}\nScore: {score}')
print('----truth----')
show_sample(random_audio)

下载 TF Lite 模型

现在您可以将 TF Lite 模型部署到您的移动或嵌入式设备。您不需要下载标签文件，因为您可以从 .tflite 文件元数据中检索标签，如前面的推理示例所示。

try:
  from google.colab import files
except ImportError:
  pass
else:
  files.download(tflite_file)

查看我们在 Android 和 iOS 上使用 TFLite 音频模型执行推理的端到端示例应用程序。