スペクトログラムを用いた音声データの学習｜音声データを画像化してCNNで学習する手法を紹介

curl https://zenodo.org/records/1188976/files/Audio_Speech_Actors_01-24.zip?download=1 -o Audio_Speech_Actors_01-24.zip

mkdir data
unzip -q Audio_Speech_Actors_01-24.zip -d data

import glob
import torch, torchvision, torchaudio
import torch.nn as nn
import torch.nn.functional as F
import librosa
import random
import os
import matplotlib.pyplot as plt
import multiprocessing
import timm
from tqdm.notebook import tqdm

class AudioDataset(torch.utils.data.Dataset) :
  def __init__(self, path, mode, sr=16000, hop_length=512, sec=2) :
    self.files = glob.glob(os.path.join(path,"*/*"))
    self.mel = torchaudio.transforms.MelSpectrogram(
        sample_rate=sr, hop_length=512, n_fft=1024, f_max=sr/2, power=1) # 1 for magnitude, 2 for power
    self.sr = sr
    self.hop_lenght = hop_length
    self.sec = sec

    # データ拡張用の変換列などを定義しておく

    if mode == "train" :
      self.mode = 1
    elif mode == "eval" :
      self.mode = 2
    else:
      self.mode = 0
  def __len__(self) :
    return len(self.files)
  def __getitem__(self, idx) :
    filename = self.files[idx]
    y, sr = librosa.load(filename, sr=self.sr)

    yt, _ = librosa.effects.trim(y)

    # trimして不足したら先頭yにする
    if len(yt) > self.sr*self.sec : y = yt

    # 正規化する（正規化した方が今回のデータでは良さそう）
    y = librosa.util.normalize(y) 

    if self.mode == 1 or self.mode == 2 :
      start = random.randint(0, len(y)-self.sr*self.sec)
    else: 
      start = 0

    y_sel = y[start:start+self.sr*self.sec]

    label = None
    if self.mode == 1 or self.mode == 2:
      # ここで、ノイズ付加などのデータ拡張を行う
      label = int(filename.split("-")[-5])-1

    mel_spec = self.mel(torch.tensor(y_sel))

    return mel_spec.unsqueeze(0), label

EPOCHS = 20
BATCH_SIZE = 8
NUM_CLASSES = 8
cpus = multiprocessing.cpu_count()

device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)

dataset = AudioDataset('data', mode="train", sec=2)
train_dataloader = torch.utils.data.DataLoader(dataset, batch_size = BATCH_SIZE, shuffle=True, num_workers=cpus)

fig, ax = plt.subplots(2, 4, figsize=(10, 8))
for inputs, labels in train_dataloader:
  for n, x in enumerate(inputs) :
    i = n//4
    j = n%4
    ax[i][j].imshow(x[0])
    ax[i][j].set_title(labels[n])
  break
plt.show()

model = timm.create_model("efficientnet_b0", pretrained=True, in_chans=1, num_classes = NUM_CLASSES).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.005)

model.train()

losses = []
acc = []

for epoch in range(EPOCHS) :
  train_loss, train_acc = 0.0, 0.0
  with tqdm(train_dataloader) as pbar:
    pbar.set_description(f'[Epoch {epoch + 1}/{EPOCHS}]')
    for inputs, labels in pbar:
      optimizer.zero_grad()
      inputs = inputs.to(device)
      labels = labels.to(device)
      outputs = model(inputs)
      loss = criterion(outputs, labels)
      loss.backward()
      optimizer.step()
      train_loss += loss.detach().item()

      preds = outputs.argmax(axis = 1)
      train_acc += torch.sum(preds == labels).item() / len(labels)

      pbar.set_postfix({'loss': loss.item(), "acc":torch.sum(preds == labels).item() / len(labels)})

  train_loss /= len(train_dataloader)
  train_acc /=  len(train_dataloader)
  losses.append(train_loss)
  acc.append(train_acc)
  print(f"epoch {epoch+1}: loss = {train_loss:.4f} acc={train_acc:.4f}")

plt.plot(losses, 'ro--')
plt.show()

plt.plot(acc, 'ro--')
plt.show()

test_dataset = AudioDataset('data', mode="eval", sec=2)
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size = BATCH_SIZE, shuffle=False, num_workers=cpus)

gt = []
pred = []
model.eval()
with tqdm(train_dataloader) as pbar:
  pbar.set_description(f'[Epoch {epoch + 1}/{EPOCHS}]')
  for inputs, labels in pbar:
    inputs = inputs.to(device)
    labels = labels.to(device)
    with torch.no_grad() :
      outputs = model(inputs)
    p = outputs.argmax(axis = 1)
    # print(p, labels)
    pbar.set_postfix({"acc":torch.sum(p == labels).item() / len(labels)})
    pred += p.tolist()
    gt += labels.tolist()

print("acc = ", sum([p == g for p, g in zip(pred, gt)]) / len(pred))

from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

cm = confusion_matrix(pred, gt)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=[i for i in range(8)])
disp.plot()