Merge branch 'develop' of https://github.com/juanmc2005/repr-learning into develop

Author: juan coria

datasets/voxceleb.py

@@ -37,26 +37,22 @@ def __next__(self):
         return batch, torch.Tensor(dic['y']).long()
 
 
-class VoxCeleb1(SimDataset):
-    sample_rate = 16000
-
-    @staticmethod
-    def config(segment_size_s: float):
-        return metrics.SpeakerValidationConfig(protocol_name='VoxCeleb.SpeakerVerification.VoxCeleb1_X',
-                                               feature_extraction=RawAudio(sample_rate=VoxCeleb1.sample_rate),
-                                               preprocessors={'audio': FileFinder()},
-                                               duration=segment_size_s)
+class VoxCelebDataset(SimDataset):
 
     def __init__(self, batch_size: int, segment_size_millis: int):
+        self.sample_rate = 16000
         self.batch_size = batch_size
         self.segment_size_s = segment_size_millis / 1000
-        self.nfeat = VoxCeleb1.sample_rate * segment_size_millis // 1000
+        self.nfeat = self.sample_rate * segment_size_millis // 1000
         self.config = VoxCeleb1.config(self.segment_size_s)
         self.protocol = get_protocol(self.config.protocol_name, preprocessors=self.config.preprocessors)
         self.train_gen, self.dev_gen, self.test_gen = None, None, None
         print(f"[Segment Size: {self.segment_size_s}s]")
         print(f"[Embedding Size: {self.nfeat}]")
 
+    def _create_config(self, segment_size_sec: float):
+        raise NotImplementedError
+
     def training_partition(self) -> VoxCelebPartition:
         if self.train_gen is None:
             self.train_gen = SpeechSegmentGenerator(self.config.feature_extraction, self.protocol,
@@ -77,3 +73,21 @@ def test_partition(self) -> VoxCelebPartition:
                                                    subset='test', per_label=1, per_fold=self.batch_size,
                                                    duration=self.segment_size_s, parallel=2)
         return VoxCelebPartition(self.test_gen, self.nfeat)
+
+
+class VoxCeleb1(VoxCelebDataset):
+
+    def _create_config(self, segment_size_sec: float):
+        return metrics.SpeakerValidationConfig(protocol_name='VoxCeleb.SpeakerVerification.VoxCeleb1_X',
+                                               feature_extraction=RawAudio(sample_rate=self.sample_rate),
+                                               preprocessors={'audio': FileFinder()},
+                                               duration=segment_size_sec)
+
+
+class VoxCeleb2(VoxCelebDataset)
+
+    def _create_config(self, segment_size_sec: float):
+        return metrics.SpeakerValidationConfig(protocol_name='VoxCeleb.SpeakerVerification.VoxCeleb2',
+                                               feature_extraction=RawAudio(sample_rate=self.sample_rate),
+                                               preprocessors={'audio': FileFinder()},
+                                               duration=segment_size_sec)

losses/config.py

@@ -100,11 +100,19 @@ def __init__(self, device, nfeat, nclass, lweight=1, distance=EuclideanDistance(
         super(CenterConfig, self).__init__('Center Loss', f"λ={lweight} - {distance}", loss_module, self.loss, distance)
 
     def optimizer(self, model, task, lr):
-        # TODO change optimizer according to task
-        # Was using lr0=0.001 and lr1=0.5
-        optimizers = [optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005),
-                      optim.SGD(self.loss.center_parameters(), lr=lr)]
-        schedulers = [lr_scheduler.StepLR(optimizers[0], 20, gamma=0.8)]
+        if task == 'mnist':
+            # Was using lr0=0.001 and lr1=0.5
+            optimizers = [optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005),
+                          optim.SGD(self.loss.center_parameters(), lr=lr)]
+            schedulers = [lr_scheduler.StepLR(optimizers[0], 20, gamma=0.8)]
+        elif task == 'speaker':
+            optimizers = sincnet_optims(model, lr)
+            optimizers.append(optim.SGD(self.loss.center_parameters(), lr=lr))
+            schedulers = []
+        elif task == 'sts':
+            optimizers, schedulers = [optim.RMSprop(model.parameters(), lr=lr)], []
+        else:
+            raise ValueError('Task must be one of mnist/speaker/sts')
         return base.Optimizer(optimizers, schedulers)
 
 
@@ -116,12 +124,20 @@ def __init__(self, device, nfeat, nclass, alpha=6.25):
         super(CocoConfig, self).__init__('CoCo Loss', f"α={alpha}", loss_module, loss, CosineDistance())
 
     def optimizer(self, model, task, lr):
-        # TODO change optimizer according to task
-        # Was using lr0=0.001 and lr1=0.01
-        params = model.all_params()
-        optimizers = [optim.SGD(params[0], lr=lr, momentum=0.9, weight_decay=0.0005),
-                      optim.SGD(params[1], lr=lr, momentum=0.9)]
-        schedulers = [lr_scheduler.StepLR(optimizers[0], 10, gamma=0.5)]
+        if task == 'mnist':
+            # Was using lr0=0.001 and lr1=0.01
+            params = model.all_params()
+            optimizers = [optim.SGD(params[0], lr=lr, momentum=0.9, weight_decay=0.0005),
+                          optim.SGD(params[1], lr=lr, momentum=0.9)]
+            schedulers = [lr_scheduler.StepLR(optimizers[0], 10, gamma=0.5)]
+        elif task == 'speaker':
+            optimizers = sincnet_optims(model, lr)
+            optimizers.append(optim.SGD(self.loss_module.parameters(), lr=lr))
+            schedulers = []
+        elif task == 'sts':
+            optimizers, schedulers = [optim.RMSprop(model.parameters(), lr=lr)], []
+        else:
+            raise ValueError('Task must be one of mnist/speaker/sts')
         return base.Optimizer(optimizers, schedulers)