I want to build wake word detection model using RNN.
I built Dataset for it.
from pathlib import Path

import pandas as pd
import torch
import torchaudio


class WakeWordDataSet(torch.utils.data.Dataset):
    def __init__(self, csv_file: Path, sample_rate: int):
        self._data = pd.read_csv(csv_file, header=None)
        self._sample_rate = sample_rate
        self._mfcc = torchaudio.transforms.MFCC(sample_rate)

    def __len__(self):
        return len(self._data)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        data = self._data.iloc[idx]
        waveform, sample_rate = torchaudio.load(data[0])
        label = data[1]
        if sample_rate > self._sample_rate:
            waveform = torchaudio.transforms.Resample(sample_rate, self._sample_rate)(waveform)
        mfcc = self._mfcc(waveform)
        return mfcc, label

It’s simply computes mfcc.
Here’s my model
import torch

import torch.nn as nn


class WakeWordDetector(nn.Module):
    def __init__(self, seq_length, hidden_size, num_layers):
        super(WakeWordDetector, self).__init__()
        self._gru = nn.GRU(seq_length, hidden_size, num_layers=num_layers, dropout=0.25)
        self._classifier = nn.Linear(num_layers * hidden_size, out_features=1)

    def forward(self, x: torch.Tensor):
        # x => batch_size, seq_len, features
        x = x.transpose(0, 1)
        _, x = self._gru(x)
        x = x.transpose(0, 1)
        x = x.flatten(1)
        # print(x.shape)
        x = self._classifier(x)
        return x.squeeze_(-1)

And here’s my training scipt
from pathlib import Path
from sklearn.metrics import f1_score
import torch
import torch.optim as optim
from torch.utils.data import DataLoader

from wakeword.dataset import WakeWordDataSet
from wakeword.model import WakeWordDetector
from wakeword.utils import collate_fn

train_data = WakeWordDataSet(Path('./wakeword/train.csv'), 8000)
test_data = WakeWordDataSet(Path('./wakeword/test.csv'), 8000)
model = WakeWordDetector(40, 64, num_layers=4)

train_loader = DataLoader(dataset=train_data, batch_size=32, collate_fn=collate_fn, shuffle=True)
test_loader = DataLoader(dataset=test_data, batch_size=32, collate_fn=collate_fn)

epochs = 100

device = 'cuda' if torch.cuda.is_available() else 'cpu'

model = model.to(device)

opt = optim.AdamW(model.parameters(), lr=0.0001)

criterion = torch.nn.BCELoss()
for epoch in range(epochs):
    running_loss = 0.0
    running_f1 = 0.0
    model.train(True)
    for i, data in enumerate(train_loader):
        mfccs = data[0].to(device)
        labels = data[1].to(device)

        opt.zero_grad()
        outputs = model(mfccs)
        outputs = torch.sigmoid(outputs)
        loss = criterion(outputs, labels)
        rounded = torch.round(outputs)
        loss.backward()
        opt.step()
        running_f1 += f1_score(labels.cpu().numpy(), rounded.cpu().detach().numpy())
        running_loss += loss.item()
        if i % 20 == 19:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 20))
            print('[%d, %5d] f1: %.3f' %
                  (epoch + 1, i + 1, running_f1 / 20))
            running_loss = 0.0
            running_f1 = 0.0

    test_f1 = 0.0
    test_loss = 0.0
    model.eval()
    with torch.no_grad():
        for i, test_data in enumerate(test_loader):
            mfccs = test_data[0].to(device)
            labels = test_data[1].to(device)

            outputs = model(mfccs)
            outputs = torch.sigmoid(outputs)
            rounded = torch.round(outputs)
            current_f1 = f1_score(labels.cpu().numpy(), rounded.cpu().detach().numpy())
            test_f1 += current_f1
            loss = criterion(outputs, labels)
            test_loss += loss.item()

        print(f'test loss: {test_loss / len(test_loader)}')
        print(f'test f1: {test_f1 / len(test_loader)}')
        if test_f1 / len(test_loader) == 1.0:
            torch.save(model.state_dict(), f'model{epoch}.pt')
            break

I have imbalanced dataset so i’m using f1 score.
So this model converges pretty fast and i can see it mysewlf using debugger. Outputs are really good. It can archive 1.0 f1 score but when i want to reload this model in some scrpit for example:
import os

import torch

import torchaudio

from wakeword.model import WakeWordDetector

checkpoint = torch.load('./model0.pt')
model = WakeWordDetector(40, 64, num_layers=4)
model.load_state_dict(checkpoint)
model.eval()
tr = torchaudio.transforms.MFCC(8000)
with torch.no_grad():
    for file in os.listdir('./sounds/train/1'):
        w, _ = torchaudio.load('./sounds/train/1/' + file)
        w = tr(w)
        w = w.transpose(1, -1)
        print((torch.sigmoid(model(w))))

it seems strange and doesn’t really work. If I use LSTM all outputs are about the same and if i use GRU outputs are slightly different. I checked inputs/outputs shapes they seem ok so I don’t know what’s the problem is:(
UPD: I’m using pytorch 1.6.0, python 3.6.11, ubuntu 20