feat(train): make num_workers configurable (#315)

voicepaw · Apr 13, 2023 · e8df714 · e8df714
1 parent 432f89b
commit e8df714
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 4 deletions.
diff --git a/src/so_vits_svc_fork/preprocessing/config_templates/quickvc.json b/src/so_vits_svc_fork/preprocessing/config_templates/quickvc.json
@@ -23,7 +23,8 @@
     "fft_sizes": [768, 1366, 342],
     "hop_sizes": [60, 120, 20],
     "win_lengths": [300, 600, 120],
-    "window": "hann_window"
+    "window": "hann_window",
+    "num_workers": 4
   },
   "data": {
     "training_files": "filelists/44k/train.txt",

diff --git a/src/so_vits_svc_fork/preprocessing/config_templates/so-vits-svc-4.0v1-legacy.json b/src/so_vits_svc_fork/preprocessing/config_templates/so-vits-svc-4.0v1-legacy.json
@@ -19,7 +19,8 @@
     "use_sr": true,
     "max_speclen": 512,
     "port": "8001",
-    "keep_ckpts": 3
+    "keep_ckpts": 3,
+    "num_workers": 4
   },
   "data": {
     "training_files": "filelists/44k/train.txt",

diff --git a/src/so_vits_svc_fork/preprocessing/config_templates/so-vits-svc-4.0v1.json b/src/so_vits_svc_fork/preprocessing/config_templates/so-vits-svc-4.0v1.json
@@ -19,7 +19,8 @@
     "use_sr": true,
     "max_speclen": 512,
     "port": "8001",
-    "keep_ckpts": 3
+    "keep_ckpts": 3,
+    "num_workers": 4
   },
   "data": {
     "training_files": "filelists/44k/train.txt",

diff --git a/src/so_vits_svc_fork/train.py b/src/so_vits_svc_fork/train.py
@@ -45,7 +45,7 @@ def train_dataloader(self):
         return DataLoader(
             self.train_dataset,
             # pin_memory=False,
-            num_workers=min(cpu_count(), 4),
+            num_workers=min(cpu_count(), self.__hparams.train.get("num_workers", 4)),
             batch_size=self.__hparams.train.batch_size,
             collate_fn=self.collate_fn,
         )