Bi-LSTM-kws.py

Bi-LSTM-kws.py

import os
import sys
import torch
import librosa
import numpy as np
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from sklearn.model_selection import train_test_split

class Net(nn.Module):
    def __init__(self, input_dim, n_classes):
        super(Net, self).__init__()
        self.l1 = nn.LSTM(input_dim, 128, 2, batch_first=True, bidirectional=True)
        self.l2 = nn.Linear(256, n_classes)
        self.l3 = nn.Softmax()
    def forward(self, x):
        out, h = self.l1(x)
        return self.l3(self.l2(out[:,-1,:]))

def prepare_data():
    X = []
    Y = []
    base = 'data'
    for kIdx, kwdir in enumerate(os.listdir(base)):
        for f in os.listdir(os.path.join(base, kwdir))[:12]:
            fpath = os.path.join(base, kwdir, f)
            y, sr = librosa.load(fpath, sr=16000, mono=True)
            """
            S2 = librosa.stft(y, n_fft=400, hop_length=160) ** 2
            Db = librosa.power_to_db(S2)
            feats = Db.T
            """
            feats = librosa.feature.melspectrogram(y, sr, n_mels=40, n_fft=400, hop_length=160, power=1)
            feats2 = librosa.feature.stack_memory(feats, n_steps=5).T
            X.append(feats2)
            Y.append(kIdx)
    return X, Y

n_epochs = 50
X, Y = prepare_data()
Xtrain, Xtest, Ytrain, Ytest = train_test_split(X, Y, test_size=0.5)
model = Net(200, 50)
optimizer = optim.Adam(model.parameters())
loss_function = nn.CrossEntropyLoss()
for epoch in range(n_epochs):
    epoch_training_loss = []
    epoch_training_acc = []
    for idx, (np_inp, np_out) in enumerate(zip(Xtrain, Ytrain)):
        inp = Variable(torch.FloatTensor(np_inp).unsqueeze(0))
        ytrue = Variable(torch.LongTensor(np.array([np_out])))
        optimizer.zero_grad()
        out = model(inp)
        loss = loss_function(out, ytrue)
        loss.backward()
        optimizer.step()
        epoch_training_loss.append(loss.data.numpy()[0])
        epoch_training_acc.append(1.0 if out.data.numpy()[0].argmax()==np_out else 0.0)
        sys.stdout.write('Training epoch (%d/%d) - loss: %f \t acc: %f \t progress: %f\r' % (epoch+1, n_epochs, np.mean(epoch_training_loss), np.mean(epoch_training_acc), float(idx+1)/len(Xtrain)))
        sys.stdout.flush()
    print ''
    epoch_testing_loss = []
    epoch_testing_acc = []
    for idx, (np_inp, np_out) in enumerate(zip(Xtest, Ytest)):
        inp = Variable(torch.FloatTensor(np_inp).unsqueeze(0))
        ytrue = Variable(torch.LongTensor(np.array([np_out])))
        out = model(inp)
        loss = loss_function(out, ytrue)
        epoch_testing_loss.append(loss.data.numpy()[0])
        epoch_testing_acc.append(1.0 if out.data.numpy()[0].argmax()==np_out else 0.0)
        sys.stdout.write('+Testing epoch (%d/%d) - loss: %f \t acc: %f \t progress: %f\r' % (epoch+1, n_epochs, np.mean(epoch_testing_loss), np.mean(epoch_testing_acc), float(idx+1)/len(Xtest)))
        sys.stdout.flush()
    print ''