Switch to using pytorch/audio

Author: ksanjeevan

README.md

@@ -24,7 +24,7 @@ Example results:
 
 - [soundfile](https://pypi.org/project/SoundFile/): audio loading
 - [torchparse](https://github.com/ksanjeevan/torchparse): .cfg easy model definition
-- [torchaudio_contrib](https://github.com/keunwoochoi/torchaudio-contrib): Audio transforms on GPU 
+- [pytorch/audio](https://github.com/pytorch/audio): Audio transforms 
 
 
 #### Features
@@ -160,6 +160,7 @@ Per fold metrics CRNN(Bidirectional, Dropout):
 - [x] CRNN entirely defined in .cfg
 - [x] Some bug in 'infer'
 - [x] Run 10-fold Cross Validation
+- [x] Switch over to pytorch/audio since the merge 
 - [ ] Comment things
 
 

data/data_manager.py

@@ -1,5 +1,5 @@
 
-import os, cv2
+import os
 import pandas as pd
 import numpy as np
 

data/data_sets.py

@@ -8,8 +8,6 @@
 import numpy as np
 import soundfile as sf
 import torch.utils.data as data
-import cv2
-
 
 class FolderDataset(data.Dataset):
 

net/audio.py

@@ -1,20 +1,128 @@
+'''
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.distributions import Normal, Uniform, HalfNormal
 
 from torchaudio_contrib import STFT, TimeStretch, MelFilterbank, ComplexNorm, ApplyFilterbank
+'''
 
 
+from torchaudio.transforms import Spectrogram, MelSpectrogram , ComplexNorm
+
+def _num_stft_bins(lengths, fft_length, hop_length, pad):
+    return (lengths + 2 * pad - fft_length + hop_length) // hop_length
+
+class MelspectrogramStretch(MelSpectrogram):
+
+    def __init__(self, hop_length=None, 
+                       sample_rate=44100, 
+                       num_mels=128, 
+                       fft_length=2048, 
+                       norm='whiten', 
+                       stretch_param=[0.4, 0.4]):
+
+        super(MelspectrogramStretch, self).__init__(sample_rate=sample_rate, 
+                                                    n_fft=fft_length, 
+                                                    hop_length=hop_length, 
+                                                    n_mels=num_mels)
+        
+        self.stft = Spectrogram(n_fft=self.n_fft, win_length=self.win_length,
+                                       hop_length=self.hop_length, pad=self.pad, 
+                                       power=None, normalized=False)
+
+        # Augmentation
+        self.prob = stretch_param[0]
+        self.random_stretch = RandomTimeStretch(stretch_param[1], 
+                                                self.hop_length, 
+                                                self.n_fft//2+1, 
+                                                fixed_rate=None)
+        
+        # Normalization (pot spec processing)
+        self.complex_norm = ComplexNorm(power=2.)
+        self.norm = SpecNormalization(norm)
+
+    def forward(self, x, lengths=None):
+        x = self.stft(x)
+
+        if lengths is not None:
+            lengths = _num_stft_bins(lengths, self.n_fft, self.hop_length, self.n_fft//2)
+            lengths = lengths.long()
+        
+        if torch.rand(1)[0] <= self.prob and self.training:
+            # Stretch spectrogram in time using Phase Vocoder
+            x, rate = self.random_stretch(x)
+            # Modify the rate accordingly
+            lengths = (lengths.float()/rate).long()+1
+        
+        x = self.complex_norm(x)
+        x = self.mel_scale(x)
+
+        # Normalize melspectrogram
+        x = self.norm(x)
+
+        if lengths is not None:
+            return x, lengths        
+        return x
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from torchaudio.transforms import TimeStretch, AmplitudeToDB 
+from torch.distributions import Uniform
+
+class RandomTimeStretch(TimeStretch):
+
+    def __init__(self, max_perc, hop_length=None, n_freq=201, fixed_rate=None):
+
+        super(RandomTimeStretch, self).__init__(hop_length, n_freq, fixed_rate)
+        self._dist = Uniform(1.-max_perc, 1+max_perc)
+
+    def forward(self, x):
+        rate = self._dist.sample().item()
+        return super(RandomTimeStretch, self).forward(x, rate), rate
+
+
+class SpecNormalization(nn.Module):
+
+    def __init__(self, norm_type, top_db=80.0):
+
+        super(SpecNormalization, self).__init__()
+
+        if 'db' == norm_type:
+            self._norm = AmplitudeToDB(stype='power', top_db=top_db)
+        elif 'whiten' == norm_type:
+            self._norm = lambda x: self.z_transform(x)
+        else:
+            self._norm = lambda x: x
+        
+    
+    def z_transform(self, x):
+        # Independent mean, std per batch
+        non_batch_inds = [1, 2, 3]
+        mean = x.mean(non_batch_inds, keepdim=True)
+        std = x.std(non_batch_inds, keepdim=True)
+        x = (x - mean)/std 
+        return x
+
+    def forward(self, x):
+        return self._norm(x)
+
+
+'''
 def amplitude_to_db(spec, ref=1.0, amin=1e-10, top_db=80):
     """
     Amplitude spectrogram to the db scale
     """
     power = spec**2
     return power_to_db(power, ref, amin, top_db)
 
-
 def power_to_db(spec, ref=1.0, amin=1e-10, top_db=80.0):
     """
     Power spectrogram to the db scale
@@ -41,7 +149,6 @@ def power_to_db(spec, ref=1.0, amin=1e-10, top_db=80.0):
     #log_spec /= log_spec.max()
     return log_spec
     
-
 def spec_whiten(spec, eps=1):    
     
     along_dim = lambda f, x: f(x, dim=-1).view(-1,1,1,1)
@@ -58,10 +165,6 @@ def spec_whiten(spec, eps=1):
     return resu
 
 
-def _num_stft_bins(lengths, fft_length, hop_length, pad):
-    return (lengths + 2 * pad - fft_length + hop_length) // hop_length
-
-
 class MelspectrogramStretch(nn.Module):
 
     def __init__(self, hop_length=None, num_mels=128, fft_length=2048, norm='whiten', stretch_param=[0.4, 0.4]):
@@ -89,12 +192,15 @@ def __init__(self, hop_length=None, num_mels=128, fft_length=2048, norm='whiten'
 
         self.counter = 0
 
+
+
     def forward(self, x, lengths=None):
         x = self.stft(x)
 
         if lengths is not None:
             lengths = _num_stft_bins(lengths, self.fft_length, self.hop_length, self.fft_length//2)
-
+            lengths = lengths.long()
+            
         if torch.rand(1)[0] <= self.prob and self.training:
             rate = 1 - self.dist.sample()
             x = self.pv(x, rate)
@@ -114,3 +220,4 @@ def __repr__(self):
         param_str = '(num_mels={}, fft_length={}, norm={}, stretch_param={})'.format(
                         self.num_mels, self.fft_length, self.norm.__name__, self.stretch_param)
         return self.__class__.__name__ + param_str
+'''

net/model.py

@@ -26,7 +26,7 @@ def __init__(self, classes, config={}, state_dict=None):
                                 stretch_param=[0.4, 0.4])
 
         # shape -> (channel, freq, token_time)
-        self.net = parse_cfg(config['cfg'], in_shape=[in_chan, self.spec.num_mels, 400])
+        self.net = parse_cfg(config['cfg'], in_shape=[in_chan, self.spec.n_mels, 400])
 
     def _many_to_one(self, t, lengths):
         return t[torch.arange(t.size(0)), lengths - 1]
@@ -39,14 +39,13 @@ def safe_param(elem):
             #if name.startswith(('conv2d','maxpool2d')):
             if isinstance(layer, (nn.Conv2d, nn.MaxPool2d)):
                 p, k, s = map(safe_param, [layer.padding, layer.kernel_size,layer.stride]) 
-                lengths = (lengths + 2*p - k)//s + 1
+                lengths = ((lengths + 2*p - k)//s + 1).long()
 
         return torch.where(lengths > 0, lengths, torch.tensor(1, device=lengths.device))
 
     def forward(self, batch):    
         # x-> (batch, time, channel)
         x, lengths, _ = batch # unpacking seqs, lengths and srs
-
         # x-> (batch, channel, time)
         xt = x.float().transpose(1,2)
         # xt -> (batch, channel, freq, time)

run.py

@@ -85,7 +85,6 @@ def train_main(config, resume):
     model = getattr(net_module, m_name)(classes, config=config)
     num_classes = len(classes)
 
-    print(model)
 
     loss = getattr(net_module, config['train']['loss'])
     metrics = getattr(net_module, config['metrics'])(num_classes)

train/base_trainer.py

@@ -18,6 +18,7 @@ def __init__(self, model, loss, metrics, optimizer, resume, config, train_logger
         self.logger = logging.getLogger(self.__class__.__name__)
 
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
         self.model = model.to(self.device)
 
         self.loss = loss