wassname/TTS
1
0
Fork 0
mirror of https://github.com/wassname/TTS.git synced 2026-09-11 12:00:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

use native amp for tacotron training

Browse Source
This commit is contained in:
erogol
2020-11-14 12:59:28 +01:00
parent 6cc464ead6
commit d8511efa8f
1 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TTS/bin/train_tacotron.py
+7 -3
View File
@@ -539,12 +539,16 @@ def main(args): # pylint: disable=redefined-outer-name
if args.restore_path:
checkpoint = torch.load(args.restore_path, map_location='cpu')
try:
# TODO: fix optimizer init, model.cuda() needs to be called before
print(" > Restoring Model.")
model.load_state_dict(checkpoint['model'])
# optimizer restore
# optimizer.load_state_dict(checkpoint['optimizer'])
print(" > Restoring Optimizer.")
optimizer.load_state_dict(checkpoint['optimizer'])
if "scaler" in checkpoint and c.mixed_precision:
print(" > Restoring AMP Scaler...")
scaler.load_state_dict(checkpoint["scaler"])
if c.reinit_layers:
raise RuntimeError
model.load_state_dict(checkpoint['model'])
except KeyError:
print(" > Partial model initialization.")
model_dict = model.state_dict()