rewrote demucs model

changed configs default values
refactoring

app.py

@@ -24,15 +24,13 @@ def denoising_transform(audio):
     )
     model.predict(audio, tgt_path)
     return tgt_path
-    # model.predict(src_path, tgt_path)
-    # return tgt_path
+
 
 
 inputs = gr.inputs.Audio(label="Source Audio", source="microphone", type='filepath')
 outputs = gr.outputs.Audio(label="Target Audio", type='filepath')
 
 title = "Denoising"
-#"""
 gr.Interface(
     denoising_transform, inputs, outputs, title=title,
     allow_flagging='never'
@@ -40,6 +38,4 @@ gr.Interface(
     server_name='localhost',
     server_port=7871,
     share=True
-    #ssl_keyfile='example.key',
-    #ssl_certfile="example.crt",
 )

conf/config.yaml

@@ -1,11 +1,11 @@
 defaults:
   - dataset: valentini
   - model: demucs
-  - loss: mse
+  - loss: L1_Multi_STFT
   - optimizer: sgd
 
 training:
-  num_epochs: 5
+  num_epochs: 10
   model_save_path: /media/public/checkpoints
 
 dataloader:
@@ -13,6 +13,8 @@ dataloader:
   sample_rate: 16000
   train_batch_size: 12
   valid_batch_size: 12
+  normalize: False
+  num_workers: 12
 
 validation:
   path: /media/public/datasets/denoising/DS_10283_2791/noisy_testset_wav
@@ -20,13 +22,13 @@ validation:
 
 
 wandb:
+  entity: borisovmaksim
   run_name: default
   project: denoising
   log_interval: 100
   api_key: local-e23d01ece807cb31e69b2cf4137e4998e4b9856f
   host: http://localhost:8080/
   notes: "Experiment note"
-  tags:
-    - baseline
+  resume: False
 
-gpu: 0
+gpu: 1

conf/model/demucs.yaml

@@ -1,5 +1,6 @@
 demucs:
  H: 64
+ L: 3
  encoder:
    conv1:
      kernel_size: 8

datasets/__init__.py

@@ -9,10 +9,9 @@ DATASETS_POOL = {
 }
 
 
-
 def get_datasets(cfg) -> Tuple[Dataset, Dataset]:
     name, dataset_params = list(cfg['dataset'].items())[0]
-    transform = Transform(input_sr=dataset_params['sample_rate'], **cfg['dataloader'])
+    transform = Transform(input_sample_rate=dataset_params['sample_rate'], **cfg['dataloader'])
     train_dataset = DATASETS_POOL[name](valid=False, transform=transform, **dataset_params)
     valid_dataset = DATASETS_POOL[name](valid=True, transform=transform, **dataset_params)
     return train_dataset, valid_dataset

denoisers/SpectralGating.py

@@ -18,6 +18,8 @@ class SpectralGating(torch.nn.Module):
         torchaudio.save(out_path, reduced_noise, rate)
         return out_path
 
+
+
     
     
 

denoisers/demucs.py

@@ -1,6 +1,7 @@
 import torch
 from torch.nn.functional import pad
-
+from utils import pad_cut_batch_audio
+import torch.nn as nn
 
 
 class Encoder(torch.nn.Module):
@@ -25,9 +26,9 @@ class Encoder(torch.nn.Module):
 
 
 class Decoder(torch.nn.Module):
-    def __init__(self, in_channels, out_channels, cfg):
+    def __init__(self, in_channels, out_channels, cfg, is_last=False):
         super(Decoder, self).__init__()
-
+        self.is_last = is_last
         self.conv1 = torch.nn.Conv1d(in_channels=in_channels, out_channels=2 * in_channels,
                                      kernel_size=cfg['conv1']['kernel_size'],
                                      stride=cfg['conv1']['stride'])
@@ -39,45 +40,46 @@ class Decoder(torch.nn.Module):
 
     def forward(self, x):
         x = self.glu(self.conv1(x))
-        x = self.relu(self.conv2(x))
+        x = self.conv2(x)
+        if not self.is_last:
+            x = self.relu(x)
         return x
 
 
 class Demucs(torch.nn.Module):
     def __init__(self, cfg):
         super(Demucs, self).__init__()
-        H = cfg['H']
-
-        self.encoder1 = Encoder(in_channels=1, out_channels=H, cfg=cfg['encoder'])
-        self.encoder2 = Encoder(in_channels=H, out_channels=2*H, cfg=cfg['encoder'])
-        self.encoder3 = Encoder(in_channels=2*H, out_channels=4*H, cfg=cfg['encoder'])
+        self.L = cfg['L']
+
+        encoders = [Encoder(in_channels=1, out_channels=cfg['H'], cfg=cfg['encoder'])]
+        decoders = [Decoder(in_channels=cfg['H'], out_channels=1, cfg=cfg['decoder'], is_last=True)]
+        for i in range(self.L - 1):
+            encoders.append(Encoder(in_channels=(2 ** i) * cfg['H'],
+                                    out_channels=(2 ** (i + 1)) * cfg['H'],
+                                    cfg=cfg['encoder']))
+            decoders.append(Decoder(in_channels=(2 ** (i + 1)) * cfg['H'],
+                                    out_channels=(2 ** i) * cfg['H'],
+                                    cfg=cfg['decoder']))
+        self.encoders = nn.ModuleList(encoders)
+        self.decoders = nn.ModuleList(decoders)
 
         self.lstm = torch.nn.LSTM(
-                                  input_size=4*H,
-                                  hidden_size=4*H, num_layers=2, batch_first=True)
-
-        self.decoder1 = Decoder(in_channels=4*H, out_channels=2*H, cfg=cfg['decoder'])
-        self.decoder2 = Decoder(in_channels=2*H, out_channels=H, cfg=cfg['decoder'])
-        self.decoder3 = Decoder(in_channels=H, out_channels=1, cfg=cfg['decoder'])
+            input_size=(2 ** (self.L - 1)) * cfg['H'],
+            hidden_size=(2 ** (self.L - 1)) * cfg['H'], num_layers=2, batch_first=True)
 
     def forward(self, x):
-        out1 = self.encoder1(x)
-        out2 = self.encoder2(out1)
-        out3 = self.encoder3(out2)
+        outs = [x]
+        for i in range(self.L):
+            out = self.encoders[i](outs[-1])
+            outs.append(out)
+        model_input = outs.pop(0)
 
-        x, _ = self.lstm(out3.permute(0, 2, 1))
+        x, _ = self.lstm(outs[-1].permute(0, 2, 1))
         x = x.permute(0, 2, 1)
-        x = self.decoder1(x + out3)
-        x = x[:, :, :out2.shape[-1]]
-        x = self.decoder2(x + out2)
-        x = x[:, :, :-1]
-        out1 = out1[:, :, :-1]
-        if x.shape[-1] > out1.shape[-1]:
-            x = x[:, :, :out1.shape[-1]]
-        elif x.shape[-1] < out1.shape[-1]:
-            out1 = out1[:, :, :x.shape[-1]]
 
-        x = self.decoder3(x + out1)
+        for i in reversed(range(self.L)):
+            decoder = self.decoders[i]
+            x = pad_cut_batch_audio(x, outs[i].shape)
+            x = decoder(x + outs[i])
+        x = pad_cut_batch_audio(x, model_input.shape)
         return x
-
-

losses.py

@@ -107,9 +107,9 @@ class MultiResolutionSTFTLoss(torch.nn.Module):
     """Multi resolution STFT loss module."""
 
     def __init__(self,
-                 fft_sizes=[1024, 2048, 512],
-                 hop_sizes=[120, 240, 50],
-                 win_lengths=[600, 1200, 240],
+                 fft_sizes=(1024, 2048, 512),
+                 hop_sizes=(120, 240, 50),
+                 win_lengths=(600, 1200, 240),
                  window="hann_window", factor_sc=0.1, factor_mag=0.1):
         """Initialize Multi resolution STFT loss module.
         Args:

main.py

@@ -3,7 +3,6 @@ from omegaconf import DictConfig
 from train import train
 
 
-
 @hydra.main(version_base=None, config_path="conf", config_name="config")
 def main(cfg: DictConfig):
     train(cfg)

testing/metrics.py

@@ -1,19 +1,27 @@
 import pesq
 from torchmetrics.audio.pesq import PerceptualEvaluationSpeechQuality
 from torchmetrics.audio.stoi import ShortTimeObjectiveIntelligibility
+from torchaudio.transforms import Resample
 import torch
 import torchaudio
 from torchmetrics import SignalNoiseRatio
 
-class Metrics:
-    def __init__(self, rate=16000):
-        self.nb_pesq = PerceptualEvaluationSpeechQuality(rate, 'wb')
-        self.stoi = ShortTimeObjectiveIntelligibility(rate, False)
+class Metrics(torch.nn.Module):
+    def __init__(self, source_rate, target_rate=16000, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.source_rate = source_rate
+        self.target_rate = target_rate
+        self.resampler = Resample(orig_freq=source_rate, new_freq=target_rate)
+        self.nb_pesq = PerceptualEvaluationSpeechQuality(target_rate, 'wb')
+        self.stoi = ShortTimeObjectiveIntelligibility(target_rate, False)
         self.snr = SignalNoiseRatio()
         
-    def calculate(self, denoised, clean):
+    def forward(self, denoised, clean):
         pesq_scores, stoi_scores = 0, 0
         for denoised_wav, clean_wav in zip(denoised, clean):
+            if self.source_rate != self.target_rate:
+                denoised_wav = self.resampler(denoised_wav)
+                clean_wav = self.resampler(clean_wav)
             try:
                 pesq_scores += self.nb_pesq(denoised_wav, clean_wav).item()
                 stoi_scores += self.stoi(denoised_wav, clean_wav).item()
@@ -22,7 +30,6 @@ class Metrics:
             except ValueError as e:
                 print(e)
 
-
         return {'PESQ': pesq_scores,
                 'STOI': stoi_scores}