Audio Dataset

util/module_mixins.py

@@ -103,9 +103,9 @@ class BaseValMixin:
             keys = list(output[0].keys())
             ident = '' if output_idx == 0 else '_train'
             summary_dict.update({f'mean{ident}_{key}': torch.mean(torch.stack([output[key]
-                                                                                      for output in output]))
-                                        for key in keys if 'loss' in key}
-                                       )
+                                                                               for output in output]))
+                                 for key in keys if 'loss' in key}
+                                )
 
             # UnweightedAverageRecall
             y_true = torch.cat([output['batch_y'] for output in output]) .cpu().numpy()
@@ -121,7 +121,45 @@ class BaseValMixin:
                 self.log(key, summary_dict[key])
 
 
-class BinaryMaskDatasetMixin:
+class BaseTestMixin:
+
+    absolute_loss = nn.L1Loss()
+    nll_loss = nn.NLLLoss()
+    bce_loss = nn.BCELoss()
+
+    def test_step(self, batch_xy, batch_idx, dataloader_idx, *args, **kwargs):
+        assert isinstance(self, LightningBaseModule)
+        batch_x, batch_y = batch_xy
+        y = self(batch_x).main_out
+        test_bce_loss = self.bce_loss(y.squeeze(), batch_y)
+        return dict(test_bce_loss=test_bce_loss,
+                    batch_idx=batch_idx, y=y, batch_y=batch_y)
+
+    def test_epoch_end(self, outputs, *_, **__):
+        assert isinstance(self, LightningBaseModule)
+        summary_dict = dict()
+
+        keys = list(outputs[0].keys())
+        summary_dict.update({f'mean_{key}': torch.mean(torch.stack([output[key]
+                                                                    for output in outputs]))
+                             for key in keys if 'loss' in key}
+                            )
+
+        # UnweightedAverageRecall
+        y_true = torch.cat([output['batch_y'] for output in outputs]) .cpu().numpy()
+        y_pred = torch.cat([output['y'] for output in outputs]).squeeze().cpu().numpy()
+
+        y_pred = (y_pred >= 0.5).astype(np.float32)
+
+        uar_score = sklearn.metrics.recall_score(y_true, y_pred, labels=[0, 1], average='macro',
+                                                 sample_weight=None, zero_division='warn')
+        uar_score = torch.as_tensor(uar_score)
+        summary_dict.update({f'uar_score': uar_score})
+        for key in summary_dict.keys():
+            self.log(key, summary_dict[key])
+
+
+class DatasetMixin:
 
     def build_dataset(self):
         assert isinstance(self, LightningBaseModule)
@@ -159,21 +197,20 @@ class BinaryMaskDatasetMixin:
             util_transforms])
 
         # Datasets
-        from datasets.binar_masks import BinaryMasksDataset
         dataset = Namespace(
             **dict(
                 # TRAIN DATASET
-                train_dataset=BinaryMasksDataset(self.params.root, setting=V.DATA_OPTIONS.train,
+                train_dataset=self.dataset_class(self.params.root, setting=V.DATA_OPTIONS.train,
                                                  use_preprocessed=self.params.use_preprocessed,
                                                  stretch_dataset=self.params.stretch,
                                                  mel_transforms=mel_transforms_train, transforms=aug_transforms),
                 # VALIDATION DATASET
-                val_train_dataset=BinaryMasksDataset(self.params.root, setting=V.DATA_OPTIONS.train,
+                val_train_dataset=self.dataset_class(self.params.root, setting=V.DATA_OPTIONS.train,
                                                      mel_transforms=mel_transforms, transforms=util_transforms),
-                val_dataset=BinaryMasksDataset(self.params.root, setting=V.DATA_OPTIONS.devel,
+                val_dataset=self.dataset_class(self.params.root, setting=V.DATA_OPTIONS.devel,
                                                mel_transforms=mel_transforms, transforms=util_transforms),
                 # TEST DATASET
-                test_dataset=BinaryMasksDataset(self.params.root, setting=V.DATA_OPTIONS.test,
+                test_dataset=self.dataset_class(self.params.root, setting=V.DATA_OPTIONS.test,
                                                 mel_transforms=mel_transforms, transforms=util_transforms),
             )
         )
@@ -190,22 +227,23 @@ class BaseDataloadersMixin(ABC):
         # sampler = RandomSampler(self.dataset.train_dataset, True, len(self.dataset.train_dataset))
         sampler = None
         return DataLoader(dataset=self.dataset.train_dataset, shuffle=True if not sampler else None, sampler=sampler,
-                          batch_size=self.params.batch_size,
+                          batch_size=self.params.batch_size, pin_memory=True,
                           num_workers=self.params.worker)
 
     # Test Dataloader
     def test_dataloader(self):
         assert isinstance(self, LightningBaseModule)
         return DataLoader(dataset=self.dataset.test_dataset, shuffle=False,
-                          batch_size=self.params.batch_size,
+                          batch_size=self.params.batch_size, pin_memory=True,
                           num_workers=self.params.worker)
 
     # Validation Dataloader
     def val_dataloader(self):
         assert isinstance(self, LightningBaseModule)
-        val_dataloader = DataLoader(dataset=self.dataset.val_dataset, shuffle=False,
+        val_dataloader = DataLoader(dataset=self.dataset.val_dataset, shuffle=False, pin_memory=True,
                                     batch_size=self.params.batch_size, num_workers=self.params.worker)
 
         train_dataloader = DataLoader(self.dataset.val_train_dataset, num_workers=self.params.worker,
+                                      pin_memory=True,
                                       batch_size=self.params.batch_size, shuffle=False)
         return [val_dataloader, train_dataloader]