modif sampler

nnet/xsets.py

@@ -85,8 +85,7 @@ class SideSampler(torch.utils.data.Sampler):
         self.num_process = num_process
         self.num_replicas = num_replicas
 
-        assert batch_size % examples_per_speaker == 0
-        #assert batch_size % (examples_per_speaker * self.num_replicas) == 0
+        assert batch_size % (examples_per_speaker * self.num_replicas) == 0
         assert (self.samples_per_speaker * self.spk_count * self.examples_per_speaker) % self.num_process == 0
 
         self.batch_size = batch_size // (self.examples_per_speaker * self.num_replicas)
@@ -148,15 +147,15 @@ class SideSampler(torch.utils.data.Sampler):
             self.index_iterator[idx] = self.labels_to_indices[value][self.segment_cursors[value]]
             self.segment_cursors[value] += 1
 
-        #self.index_iterator = numpy.repeat(self.index_iterator, self.num_replicas)
-        self.index_iterator = self.index_iterator.reshape(-1, self.num_process * self.examples_per_speaker)[:, self.rank * self.examples_per_speaker:(self.rank + 1) * self.examples_per_speaker].flatten()
-        #self.index_iterator = self.index_iterator.reshape(-1, self.num_process * self.examples_per_speaker * self.num_replicas)[:, self.rank * self.examples_per_speaker * self.num_replicas:(self.rank + 1) * self.examples_per_speaker * self.num_replicas].flatten()
+        self.index_iterator = numpy.repeat(self.index_iterator, self.num_replicas)
+        #self.index_iterator = self.index_iterator.reshape(-1, self.num_process * self.examples_per_speaker)[:, self.rank * self.examples_per_speaker:(self.rank + 1) * self.examples_per_speaker].flatten()
+        self.index_iterator = self.index_iterator.reshape(-1, self.num_process * self.examples_per_speaker * self.num_replicas)[:, self.rank * self.examples_per_speaker * self.num_replicas:(self.rank + 1) * self.examples_per_speaker * self.num_replicas].flatten()
 
         return iter(self.index_iterator)
 
     def __len__(self) -> int:
-        #return (self.samples_per_speaker * self.spk_count * self.examples_per_speaker * self.num_replicas) // self.num_process
-        return (self.samples_per_speaker * self.spk_count * self.examples_per_speaker) // self.num_process
+        return (self.samples_per_speaker * self.spk_count * self.examples_per_speaker * self.num_replicas) // self.num_process
+        #return (self.samples_per_speaker * self.spk_count * self.examples_per_speaker) // self.num_process
 
     def set_epoch(self, epoch: int) -> None:
         self.epoch = epoch