Clasificación de audio utilizando embeddings de Wav2vec

Wav2Vec2 fue propuesto en el paper wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations por Alexei Baevski, Henry Zhou, Abdelrahman Mohamed, y Michael Auli.

Preparacion del ambiente

Instalamos las librerias necesarias: librosa, transformers, datasets, evaluate.

[1]:

%pip install transformers[torch] datasets accelerate evaluate --quiet

     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 7.2/7.2 MB 57.8 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 486.2/486.2 kB 34.8 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 244.2/244.2 kB 21.5 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 81.4/81.4 kB 9.8 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 268.8/268.8 kB 30.2 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 7.8/7.8 MB 79.9 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.3/1.3 MB 47.3 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 110.5/110.5 kB 15.2 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 212.5/212.5 kB 26.1 MB/s eta 0:00:00
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 134.3/134.3 kB 18.2 MB/s eta 0:00:00

Descargaremos un conjunto de datos de entrenamiento para clasificar sonidos de diferentes animales.

Nota: La siguiente celda realiza varias manipulaciones de directorios para que el nombre del directorio coincida con el nombre para mostrar de la clase.

[12]:

!git clone https://github.com/YashNita/Animal-Sound-Dataset
!mv Animal-Sound-Dataset/Aslan Animal-Sound-Dataset/Leon
!mv Animal-Sound-Dataset/Esek Animal-Sound-Dataset/Burro
!mv Animal-Sound-Dataset/Inek Animal-Sound-Dataset/Vaca
!mv Animal-Sound-Dataset/Kedi-Part1 Animal-Sound-Dataset/Gato
!mv Animal-Sound-Dataset/Kedi-Part2/* Animal-Sound-Dataset/Gato
!rm -d Animal-Sound-Dataset/Kedi-Part2
!mv Animal-Sound-Dataset/Kopek-Part1 Animal-Sound-Dataset/Perro
!mv Animal-Sound-Dataset/Kopek-Part2/* Animal-Sound-Dataset/Perro
!rm -d Animal-Sound-Dataset/Kopek-Part2
!mv Animal-Sound-Dataset/Koyun Animal-Sound-Dataset/Oveja
!mv Animal-Sound-Dataset/Kus-Part1 Animal-Sound-Dataset/Pajaro
!mv Animal-Sound-Dataset/Kus-Part2/* Animal-Sound-Dataset/Pajaro
!rm -d Animal-Sound-Dataset/Kus-Part2
!mv Animal-Sound-Dataset/Maymun Animal-Sound-Dataset/Mono
!mv Animal-Sound-Dataset/Tavuk Animal-Sound-Dataset/Gallina
!mv Animal-Sound-Dataset/Kurbaga Animal-Sound-Dataset/Rana

Cloning into 'Animal-Sound-Dataset'...
remote: Enumerating objects: 887, done.
remote: Total 887 (delta 0), reused 0 (delta 0), pack-reused 887
Receiving objects: 100% (887/887), 100.68 MiB | 16.56 MiB/s, done.
Resolving deltas: 100% (69/69), done.
Updating files: 100% (876/876), done.

Este conjunto de datos dispone de sonidos de diferentes animales, donde el nombre del directorio corresponde al animal. Este tipo de conjuntos de datos se pueden cargar facilmente utilizando la libraría datasets:

[13]:

from datasets import load_dataset, Audio

dataset = load_dataset("audiofolder", data_dir="/content/Animal-Sound-Dataset", split='train')

Downloading and preparing dataset audiofolder/default to /root/.cache/huggingface/datasets/audiofolder/default-e50b2603353acda2/0.0.0/6cbdd16f8688354c63b4e2a36e1585d05de285023ee6443ffd71c4182055c0fc...

Dataset audiofolder downloaded and prepared to /root/.cache/huggingface/datasets/audiofolder/default-e50b2603353acda2/0.0.0/6cbdd16f8688354c63b4e2a36e1585d05de285023ee6443ffd71c4182055c0fc. Subsequent calls will reuse this data.

Veamos como luce el conjunto de datos:

[14]:

dataset

[14]:

Dataset({
    features: ['audio', 'label'],
    num_rows: 875
})

La columna label tiene la categoría a la que pertene el audio. Construiremos un diccionario para transformar de indices a etiquetas, el cual nos será de utilidad luego:

[15]:

labels = dataset.features["label"].names
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
    label2id[label] = i
    id2label[i] = label

El conjunto de datos dispone de las siguientes etiquetas:

[16]:

label2id

[16]:

{'Burro': 0,
 'Gallina': 1,
 'Gato': 2,
 'Leon': 3,
 'Mono': 4,
 'Oveja': 5,
 'Pajaro': 6,
 'Perro': 7,
 'Rana': 8,
 'Vaca': 9}

Separaremos el conjunto de datos en entrenamiento y testing:

[17]:

dataset = dataset.train_test_split(test_size=0.2)

[18]:

dataset

[18]:

DatasetDict({
    train: Dataset({
        features: ['audio', 'label'],
        num_rows: 700
    })
    test: Dataset({
        features: ['audio', 'label'],
        num_rows: 175
    })
})

Utilizando un modelo preentrenado de wave2vec

Trabajando con archivos de audio

De igual forma que un modelo de NLP está entrenado para trabajar sobre tokens, el modelo de wav2vec está entrenado para trabajar sobre la una onda de audio (wave). Ya hemos visto el concepto de onda en este curso, sin embargo en esta ocación, la libraría datasets nos permite convertir una columna que tiene la ubicación de un archivo de audio en una columna de tipo Audio con la información de la onda en la misma.

Para esto, utilizaremos el método cast():

[20]:

dataset = dataset.cast_column("audio", Audio(sampling_rate=16_000))
dataset["train"][0]

[20]:

{'audio': {'path': '/content/Animal-Sound-Dataset/Vaca/inek_12.wav',
  'array': array([-0.02970379, -0.0435513 , -0.05235241, ..., -0.01638332,
         -0.01580101, -0.00682969]),
  'sampling_rate': 16000},
 'label': 9}

Feature extractors

Los FeatureExtractor nos permiten transformar un conjunto de datos con X features a otro con X”. De igual forma que un tokenizer, estos componentes se pueden descargar desde HuggingFace y cada modelo puede ser empaquetado con su correspondiente FeatureExtractor:

[19]:

from transformers import AutoFeatureExtractor

feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")

/usr/local/lib/python3.10/dist-packages/transformers/configuration_utils.py:380: UserWarning: Passing `gradient_checkpointing` to a config initialization is deprecated and will be removed in v5 Transformers. Using `model.gradient_checkpointing_enable()` instead, or if you are using the `Trainer` API, pass `gradient_checkpointing=True` in your `TrainingArguments`.
  warnings.warn(

En particular, este feature extractor realiza los siguientes pasos:

Carga los datos en memoria desde los archivos de audio.
Verifica el sampling rate y lo modifica para que conincida con el modelo en caso de que no.
Realiza el padding correspondiente de la secuencia de datos para que todos los lotes tengan la misma longitud.
Realiza la normalización de los datos de entrada, tal como es esperado por el modelo.

Para aplicar este FeatureExtractor a los datos, definiremos una función que realiza el procesamiento:

[21]:

def preprocess(examples):
    audio_arrays = [x["array"] for x in examples["audio"]]
    inputs = feature_extractor(
        audio_arrays, sampling_rate=feature_extractor.sampling_rate, max_length=16000, truncation=True
    )
    return inputs

Luego, utilizando la función map sobre el conjunto de datos aplicamos la transformación a todos los splits:

[22]:

encoded_data = dataset.map(preprocess, remove_columns="audio", batched=True)

Podemos verificar como lucen los datos luego:

[35]:

encoded_data

[35]:

DatasetDict({
    train: Dataset({
        features: ['label', 'input_values'],
        num_rows: 700
    })
    test: Dataset({
        features: ['label', 'input_values'],
        num_rows: 175
    })
})

[41]:

import numpy as np

print("Shape:", np.asarray(encoded_data['train']['input_values'][0]).shape)
print("Mean:", np.asarray(encoded_data['train']['input_values'][0]).mean())
print("STD:", np.asarray(encoded_data['train']['input_values'][0]).std())

Shape: (16000,)
Mean: 2.278744159411872e-08
STD: 0.9999912629930756

Clasificador

Ahora es momento de crear nuestro clasificador. De igual forma que hicimos con los Transformers para texto, aquí utilizaremos la clase AutoModelForAudioClassification la cual nos permite utilizar los embeddings de un modelo preentrenado, en este caso un modelo de wav2vec para luego concatenarle un simple clasificador (MLP) para resolver la tarea en cuestión.

En nuestro caso, utilizaremos un modelo preentrenado de Facebook. Note como configuramos la cantidad de clases a predecir y las etiquetas correspondientes:

[23]:

from transformers import AutoModelForAudioClassification, TrainingArguments, Trainer

num_labels = len(id2label)
model = AutoModelForAudioClassification.from_pretrained(
    "facebook/wav2vec2-base", num_labels=num_labels, label2id=label2id, id2label=id2label
)

/usr/local/lib/python3.10/dist-packages/transformers/configuration_utils.py:380: UserWarning: Passing `gradient_checkpointing` to a config initialization is deprecated and will be removed in v5 Transformers. Using `model.gradient_checkpointing_enable()` instead, or if you are using the `Trainer` API, pass `gradient_checkpointing=True` in your `TrainingArguments`.
  warnings.warn(

Some weights of the model checkpoint at facebook/wav2vec2-base were not used when initializing Wav2Vec2ForSequenceClassification: ['project_q.bias', 'quantizer.weight_proj.weight', 'project_q.weight', 'quantizer.weight_proj.bias', 'project_hid.weight', 'quantizer.codevectors', 'project_hid.bias']
- This IS expected if you are initializing Wav2Vec2ForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing Wav2Vec2ForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of Wav2Vec2ForSequenceClassification were not initialized from the model checkpoint at facebook/wav2vec2-base and are newly initialized: ['projector.weight', 'classifier.weight', 'classifier.bias', 'projector.bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

Nuestro problema es de clasificación, por lo cual mediremos el accuracy de nuestro modelo. Utilizando la libraria evaluate podemos cargar esta métrica:

[24]:

import evaluate

accuracy = evaluate.load("accuracy")

[25]:

import numpy as np

def compute_metrics(eval_pred):
    predictions = np.argmax(eval_pred.predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=eval_pred.label_ids)

Entrenando el modelo

En momento de iniciar el procedimiento de entrenamiento:

[26]:

import torch

training_args = TrainingArguments(
    output_dir="model",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    per_device_train_batch_size=32,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=32,
    num_train_epochs=10,
    warmup_ratio=0.1,
    load_best_model_at_end=True,
    optim="adamw_torch"
)

Utilizaremos los conjuntos de datos de entrenamiento y testing que separamos en un principio:

[27]:

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=encoded_data["train"],
    eval_dataset=encoded_data["test"],
    tokenizer=feature_extractor,
    compute_metrics=compute_metrics,
)

trainer.train()

[50/50 04:59, Epoch 9/10]

Epoch	Training Loss	Validation Loss	Accuracy
0	No log	2.205360	0.217143
2	No log	1.969920	0.502857
2	No log	1.826615	0.560000
4	No log	1.684681	0.582857
4	No log	1.584715	0.634286
6	No log	1.512387	0.657143
6	No log	1.471338	0.668571
8	No log	1.401476	0.708571
8	No log	1.375105	0.714286
9	No log	1.375185	0.714286

[27]:

TrainOutput(global_step=50, training_loss=1.7717454528808594, metrics={'train_runtime': 309.5449, 'train_samples_per_second': 22.614, 'train_steps_per_second': 0.162, 'total_flos': 5.7777674221824e+16, 'train_loss': 1.7717454528808594, 'epoch': 9.09})

Verificando el modelo

Podemos probar el modelo con un ejemplo. Tomemos un archivo aleatorio del directorio de datos:

[30]:

wav_file_name = '/content/Animal-Sound-Dataset/Pajaro/Kus_161.wav'

Podemos reproducir este archivo de audio para familiarizarnos con él:

[31]:

import librosa

sample_rate: float = 16000.0
wave = librosa.load(wav_file_name, sr=sample_rate)

[32]:

from IPython.display import Audio

Audio(wave[0], rate=sample_rate)

[32]:

Creamos un pipeline con el modelo que entrenamos anteriormente:

[33]:

from transformers import pipeline

classifier = pipeline("audio-classification", feature_extractor=feature_extractor, model=trainer.model, device=trainer.model.device)

[34]:

classifier([wav_file_name], top_k=1)

[34]:

[[{'score': 0.35037851333618164, 'label': 'Pajaro'}]]