mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-13 23:53:02 +00:00
Brandon Thomas
70 lines
2.4 KiB
Diff
70 lines
2.4 KiB
Diff
--- /content/acoustic-model/train.py 2022-09-03 22:51:08.632610312 -0700
|
|
+++ train(3).py 2022-09-03 23:15:29.700072538 -0700
|
|
@@ -17,7 +17,7 @@
|
|
from acoustic.utils import Metric, save_checkpoint, load_checkpoint, plot_spectrogram
|
|
|
|
|
|
-logging.basicConfig(level=logging.DEBUG)
|
|
+logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
########################################################################################
|
|
@@ -28,10 +28,10 @@
|
|
LEARNING_RATE = 4e-4
|
|
BETAS = (0.8, 0.99)
|
|
WEIGHT_DECAY = 1e-5
|
|
-STEPS = 80000
|
|
+STEPS = 20000
|
|
LOG_INTERVAL = 5
|
|
-VALIDATION_INTERVAL = 1000
|
|
-CHECKPOINT_INTERVAL = 1000
|
|
+VALIDATION_INTERVAL = 2000
|
|
+CHECKPOINT_INTERVAL = 2000
|
|
BACKEND = "nccl"
|
|
INIT_METHOD = "tcp://localhost:54321"
|
|
|
|
@@ -200,6 +200,19 @@
|
|
# Start validation loop
|
|
# --------------------------------------------------------------------------#
|
|
|
|
+
|
|
+
|
|
+ if global_step % CHECKPOINT_INTERVAL == 0:
|
|
+ save_checkpoint(
|
|
+ checkpoint_dir=args.checkpoint_dir,
|
|
+ acoustic=acoustic,
|
|
+ optimizer=optimizer,
|
|
+ step=global_step,
|
|
+ loss=validation_loss.value,
|
|
+ best=9999,
|
|
+ logger=logger,
|
|
+ )
|
|
+
|
|
if global_step % VALIDATION_INTERVAL == 0:
|
|
acoustic.eval()
|
|
validation_loss.reset()
|
|
@@ -242,22 +255,7 @@
|
|
f"valid -- epoch: {epoch}, loss: {validation_loss.value:.4f}"
|
|
)
|
|
|
|
- new_best = best_loss > validation_loss.value
|
|
- if new_best or global_step % CHECKPOINT_INTERVAL:
|
|
- if new_best:
|
|
- logger.info("-------- new best model found!")
|
|
- best_loss = validation_loss.value
|
|
-
|
|
- if rank == 0:
|
|
- save_checkpoint(
|
|
- checkpoint_dir=args.checkpoint_dir,
|
|
- acoustic=acoustic,
|
|
- optimizer=optimizer,
|
|
- step=global_step,
|
|
- loss=validation_loss.value,
|
|
- best=new_best,
|
|
- logger=logger,
|
|
- )
|
|
+
|
|
|
|
# -----------------------------------------------------------------------------#
|
|
# End validation loop
|