From e8df7146b0d1d3ee32af576c251f47d8fdd80bb3 Mon Sep 17 00:00:00 2001
From: 34j <55338215+34j@users.noreply.github.com>
Date: Thu, 13 Apr 2023 21:28:47 +0900
Subject: [PATCH] feat(train): make num_workers configurable (#315)

---
 .../preprocessing/config_templates/quickvc.json                | 3 ++-
 .../config_templates/so-vits-svc-4.0v1-legacy.json             | 3 ++-
 .../preprocessing/config_templates/so-vits-svc-4.0v1.json      | 3 ++-
 src/so_vits_svc_fork/train.py                                  | 2 +-
 4 files changed, 7 insertions(+), 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
index 7f1334d5..7a8f42ed 100644
--- 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
index 2d81e27c..a003ece1 100644
--- 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
index 2f2f9358..694c7ee7 100644
--- 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
index e3dce7ac..17e3039b 100644
--- 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,
         )