fix(train): add bf16 and fp16 support (#268)

voicepaw · Apr 9, 2023 · 4229fd8 · 4229fd8
1 parent 1b67c6f
commit 4229fd8
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 0 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
@@ -9,6 +9,7 @@
     "eps": 1e-9,
     "batch_size": 16,
     "fp16_run": false,
+    "bf16_run": false,
     "lr_decay": 0.999875,
     "segment_size": 10240,
     "init_lr_ratio": 1,

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
@@ -9,6 +9,7 @@
     "eps": 1e-9,
     "batch_size": 16,
     "fp16_run": false,
+    "bf16_run": false,
     "lr_decay": 0.999875,
     "segment_size": 10240,
     "init_lr_ratio": 1,

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
@@ -9,6 +9,7 @@
     "eps": 1e-9,
     "batch_size": 16,
     "fp16_run": false,
+    "bf16_run": false,
     "lr_decay": 0.999875,
     "segment_size": 10240,
     "init_lr_ratio": 1,

diff --git a/src/so_vits_svc_fork/train.py b/src/so_vits_svc_fork/train.py
@@ -75,6 +75,11 @@ def train(
             / len(datamodule.train_dataset)
             * hparams.train.batch_size
         ),
+        precision=16
+        if hparams.train.fp16_run
+        else "bf16"
+        if hparams.train.get("bf16_run", False)
+        else 32,
     )
     model = VitsLightning(reset_optimizer=reset_optimizer, **hparams)
     trainer.fit(model, datamodule=datamodule)