Sturcturing code

Author: raccooncoder

.gitignore

@@ -1,3 +1,8 @@
+*.pth
+*.wav
+speech_commands/
+wandb/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

README.md

@@ -1,2 +1,6 @@
 # kws-pytorch
 A KWS model trained on SpeechCommands dataset, written in PyTorch.
+
+W&B logs:
+
+https://wandb.ai/raccooncoder/kws-dlaudio

augmentations.py

@@ -0,0 +1,31 @@
+import torch
+from torch import nn
+import torchaudio
+
+youtube_noise, _ = torchaudio.load('Cafe sounds ~ Ambient noise-i9a6ReFTHiw.wav')
+youtube_noise = youtube_noise.sum(dim=0)
+
+class GaussianNoise(nn.Module):    
+    def __init__(self, mean=0, std=0.05):
+        super(GaussianNoise, self).__init__()
+        
+        self.noiser = torch.distributions.Normal(mean, std)   
+            
+    def forward(self, wav):
+        wav = wav + self.noiser.sample(wav.size())     
+        wav = wav.clamp(-1, 1)
+        
+        return wav
+    
+class YoutubeNoise(nn.Module):    
+    def __init__(self, alpha=0.05):
+        super(YoutubeNoise, self).__init__()
+        
+        self.alpha = alpha
+        self.noise_wav = youtube_noise
+            
+    def forward(self, wav):
+        wav = wav + self.alpha * self.noise_wav[:wav.shape[-1]] 
+        wav = wav.clamp(-1, 1)
+        
+        return wav

configs/.gitkeep renamed to checkpoints/.gitkeep

@@ -0,0 +1,25 @@
+{
+    "target_class": "marvin",
+    "num_epochs": 20,
+    "batch_size": 256,
+    "random_seed": 13,
+    "img_padding_length": 130,
+    "enc_hidden_size": 128,
+    "window_size": 100,
+    "conv_out_channels": 16,
+    "conv_kernel_size": 51,
+    "learning_rate": 0.001,
+    "dataloader_num_workers": 8,
+    "weight_decay": 0.001,
+    "lr_scheduler_step_size": 10,
+    "lr_scheduler_gamma": 0.1,
+    "melspec_sample_rate": 16000,
+    "melspec_n_mels": 64,
+    "melspec_n_fft": 512,
+    "melspec_hop_length": 128,
+    "melspec_f_max": 4000,
+    "specaug_freq_mask_param": 5,
+    "specaug_time_mask_param": 5,
+    "confidence_threshold": 0.9,
+    "teacher_alpha": 0.6
+}

configs/config.json

@@ -0,0 +1,30 @@
+import json
+
+config = dict(
+    target_class = 'marvin',
+    num_epochs = 20,
+    batch_size = 256,
+    random_seed = 13,
+    img_padding_length = 130,
+    enc_hidden_size = 128,
+    window_size = 100,
+    conv_out_channels = 16,
+    conv_kernel_size = 51, 
+    learning_rate = 1e-3,
+    dataloader_num_workers = 8,
+    weight_decay = 1e-3,
+    lr_scheduler_step_size = 10,
+    lr_scheduler_gamma = 0.1,
+    melspec_sample_rate = 16000,
+    melspec_n_mels = 64,
+    melspec_n_fft = 512,
+    melspec_hop_length = 128,
+    melspec_f_max = 4000,
+    specaug_freq_mask_param = 5,
+    specaug_time_mask_param = 5,
+    confidence_threshold = 0.9,
+    teacher_alpha = 0.6
+)
+
+with open('config.json', 'w') as f:
+    json.dump(config, f, indent=4)

configs/config_gen.py

@@ -0,0 +1,43 @@
+import torch
+from torch import nn
+import torchaudio
+from torch.utils.data import Dataset
+from augmentations import *
+
+class SpeechCommands(Dataset):
+    def __init__(self, config, X, y, train=True):
+        self.paths = X
+        self.labels = y
+        self.train = train
+        self.config = config
+        
+    def __len__(self):
+        return len(self.paths)
+    
+    def __getitem__(self, idx):
+        img = torch.zeros(1, self.config.melspec_n_mels, self.config.img_padding_length)
+        wav, sr = torchaudio.load(self.paths[idx])
+        
+        if self.train:
+            wav_proc = nn.Sequential(#GaussianNoise(0, 0.01),
+                                    YoutubeNoise(0.1),  
+                                    torchaudio.transforms.MelSpectrogram(sample_rate=self.config.melspec_sample_rate, 
+                                                                        n_mels=self.config.melspec_n_mels, 
+                                                                        n_fft=self.config.melspec_n_fft, 
+                                                                        hop_length=self.config.melspec_hop_length, 
+                                                                        f_max=self.config.melspec_f_max),
+                                    torchaudio.transforms.FrequencyMasking(freq_mask_param=self.config.specaug_freq_mask_param),
+                                    torchaudio.transforms.TimeMasking(time_mask_param=self.config.specaug_time_mask_param)
+                                    )
+        else:
+            wav_proc = nn.Sequential(torchaudio.transforms.MelSpectrogram(sample_rate=self.config.melspec_sample_rate, 
+                                                                        n_mels=self.config.melspec_n_mels, 
+                                                                        n_fft=self.config.melspec_n_fft, 
+                                                                        hop_length=self.config.melspec_hop_length, 
+                                                                        f_max=self.config.melspec_f_max),
+                                    )
+            
+        mel_spectrogram = torch.log(wav_proc(wav) + 1e-9)
+        img[0, :, :mel_spectrogram.size(2)] = mel_spectrogram
+        
+        return img.reshape(self.config.melspec_n_mels, self.config.img_padding_length), self.labels[idx]

dataset.py

@@ -0,0 +1,89 @@
+import wandb
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from sklearn.model_selection import train_test_split
+
+import json
+import glob
+import pandas as pd
+import numpy as np
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+torch.backends.cudnn.benchmark = True
+
+with open('configs/config.json', 'r') as f:
+    config = json.load(f)
+
+wandb.init(config=config, project="kws-dlaudio")
+config = wandb.config
+
+from utils import set_seed
+from train import train_distill, evaluate
+from inference import inference
+from dataset import SpeechCommands
+from model import KWSNet
+
+set_seed(config.random_seed)
+
+print(device)
+
+paths = []
+labels = []
+
+for path in glob.glob('speech_commands/*/*.wav'):
+    _, label, _ = path.split('/')
+    paths.append(path)
+    labels.append(int(label == config.target_class))
+    
+df = pd.DataFrame({'path': paths, 'label': labels})
+
+X_train, X_test, y_train, y_test = train_test_split(np.array(df['path']), 
+                                                    np.array(df['label']), 
+                                                    test_size=0.1, 
+                                                    stratify=np.array(df['label']),
+                                                    random_state=config.random_seed)
+
+train_dataset = SpeechCommands(config, X_train, y_train)
+test_dataset = SpeechCommands(config, X_test, y_test)
+
+train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True, num_workers=config.dataloader_num_workers, pin_memory=True)
+val_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False, num_workers=config.dataloader_num_workers, pin_memory=True)  
+
+student_model = KWSNet(config.enc_hidden_size // 2, config.conv_out_channels, config.conv_kernel_size)
+student_model = student_model.to(device)
+
+error = nn.CrossEntropyLoss()
+optimizer = torch.optim.Adam(student_model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay)
+
+lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=config.lr_scheduler_step_size, gamma=config.lr_scheduler_gamma)
+
+teacher_model = KWSNet(config.enc_hidden_size, config.conv_out_channels, config.conv_kernel_size)
+teacher_model.load_state_dict(torch.load('checkpoints/teacher_model.pth'))
+teacher_model = teacher_model.to(device)
+
+alpha = config.teacher_alpha   
+
+for epoch in range(config.num_epochs): 
+    train_distill(epoch, teacher_model, student_model, alpha, optimizer, error, train_loader, device)
+    evaluate(student_model, optimizer, error, val_loader, device)
+    lr_scheduler.step()
+
+negative_val = []
+positive_val = []
+
+for path, label in zip(X_test, y_test):
+    if label == 1:
+        positive_val.append(path)
+    else:
+        negative_val.append(path)
+
+
+path = positive_val[1]
+inference('results/student_positive_example.png', student_model, path, noise=True, device=device)
+
+path = negative_val[1]
+inference('results/student_negative_example.png', student_model, path, noise=True, device=device)
+
+torch.save(student_model.state_dict(), 'checkpoints/student_model.pth')
+wandb.save('checkpoints/student_model.pth')

distill.py

@@ -0,0 +1,21 @@
+FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04
+
+ENV TZ=Europe/Moscow
+ENV TERM xterm-256color
+
+RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
+
+RUN apt-get install -y \
+	python3-pip \
+	python3-tk \
+	libboost-all-dev
+
+RUN apt-get -y install git
+RUN python3 -m pip install --upgrade pip
+COPY requirements.txt .
+RUN pip3 install -r requirements.txt
+
+RUN wget http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz -O speech_commands_v0.01.tar.gz
+RUN mkdir speech_commands && tar -C speech_commands -xvzf speech_commands_v0.01.tar.gz 1> log
+COPY mv.sh .
+ENTRYPOINT mv.sh

docker/Dockerfile

@@ -0,0 +1 @@
+docker build -t kws_pytorch .

docker/build.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+mv speech_commands /home/kws-pytorch
+cd /home/kws-pytorch

docker/mv.sh

@@ -4,4 +4,5 @@ torch == 1.6.0
 torchaudio == 0.6.0
 wandb == 0.10.7
 librosa
-tqdm
+tqdm
+plotly

requirements.txt renamed to docker/requirements.txt

@@ -0,0 +1 @@
+docker run -v  $(pwd)/:/home/asr-pytorch -it kws_pytorch

docker/run.sh

@@ -0,0 +1,48 @@
+import wandb
+import torch
+import torchaudio
+from torch import nn
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+
+import time 
+
+config = wandb.config
+
+def inference(fname, model, path, noise, device):
+    model.eval() 
+    model.inference_mode()
+    
+    noise_wav1, _ = torchaudio.load('LJ001-0001.wav')
+    noise_wav2, _ = torchaudio.load('LJ001-0014.wav')
+    
+    wav, sr = torchaudio.load(path)
+    wav_proc = nn.Sequential(torchaudio.transforms.MelSpectrogram(sample_rate=config.melspec_sample_rate, 
+                                                                    n_mels=config.melspec_n_mels, 
+                                                                    n_fft=config.melspec_n_fft, 
+                                                                    hop_length=config.melspec_hop_length, 
+                                                                    f_max=config.melspec_f_max))
+    mel_spectrogram = torch.log(wav_proc(wav) + 1e-9)
+
+    if noise:
+        noise1_melspec = torch.log(wav_proc(noise_wav1) + 1e-9) 
+        noise2_melspec = torch.log(wav_proc(noise_wav2) + 1e-9)
+    
+        img = torch.cat((noise1_melspec, mel_spectrogram, noise2_melspec), -1)
+    else:
+        img = mel_spectrogram
+    
+    img = img.to(device)
+    
+    start = time.time()
+    with torch.no_grad():
+        outputs = F.softmax(model(img).squeeze(1), dim=-1).detach().cpu().numpy()[:, 0, 1]
+
+    finish = time.time()
+    wandb.log({'Inference time': finish - start})
+    
+    plt.figure(figsize=(20,10))
+    plt.plot(range(len(outputs)), outputs)
+    plt.axhline(y=config.confidence_threshold, color='r', linestyle='-')
+    wandb.log({fname: wandb.Image(plt)})
+    plt.savefig(fname, dpi=500)