fix(task): fix estimation of training set size (#1644)

pyannote · Feb 8, 2024 · c3cc56c · c3cc56c
1 parent e611d1c
commit c3cc56c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,6 +16,7 @@
 ### Fixes
 
 - fix(task): fix random generators
+- fix(task): fix estimation of training set size
 
 ### Improvements
 

diff --git a/pyannote/audio/tasks/segmentation/mixins.py b/pyannote/audio/tasks/segmentation/mixins.py
@@ -255,8 +255,11 @@ def collate_fn(self, batch, stage="train"):
 
     def train__len__(self):
         # Number of training samples in one epoch
+        train_file_ids = np.where(
+            self.prepared_data["audio-metadata"]["subset"] == Subsets.index("train")
+        )[0]
 
-        duration = np.sum(self.prepared_data["audio-annotated"])
+        duration = np.sum(self.prepared_data["audio-annotated"][train_file_ids])
         return max(self.batch_size, math.ceil(duration / self.duration))
 
     def prepare_validation(self, prepared_data: Dict):