encoder_train.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. from utils.argutils import print_args
  2. from encoder.train import train
  3. from pathlib import Path
  4. import argparse
  5. if __name__ == "__main__":
  6. parser = argparse.ArgumentParser(
  7. description="Trains the speaker encoder. You must have run encoder_preprocess.py first.",
  8. formatter_class=argparse.ArgumentDefaultsHelpFormatter
  9. )
  10. parser.add_argument("run_id", type=str, help= \
  11. "Name for this model. By default, training outputs will be stored to saved_models/<run_id>/. If a model state "
  12. "from the same run ID was previously saved, the training will restart from there. Pass -f to overwrite saved "
  13. "states and restart from scratch.")
  14. parser.add_argument("clean_data_root", type=Path, help= \
  15. "Path to the output directory of encoder_preprocess.py. If you left the default "
  16. "output directory when preprocessing, it should be <datasets_root>/SV2TTS/encoder/.")
  17. parser.add_argument("-m", "--models_dir", type=Path, default="saved_models", help=\
  18. "Path to the root directory that contains all models. A directory <run_name> will be created under this root."
  19. "It will contain the saved model weights, as well as backups of those weights and plots generated during "
  20. "training.")
  21. parser.add_argument("-v", "--vis_every", type=int, default=10, help= \
  22. "Number of steps between updates of the loss and the plots.")
  23. parser.add_argument("-u", "--umap_every", type=int, default=100, help= \
  24. "Number of steps between updates of the umap projection. Set to 0 to never update the "
  25. "projections.")
  26. parser.add_argument("-s", "--save_every", type=int, default=500, help= \
  27. "Number of steps between updates of the model on the disk. Set to 0 to never save the "
  28. "model.")
  29. parser.add_argument("-b", "--backup_every", type=int, default=7500, help= \
  30. "Number of steps between backups of the model. Set to 0 to never make backups of the "
  31. "model.")
  32. parser.add_argument("-f", "--force_restart", action="store_true", help= \
  33. "Do not load any saved model.")
  34. parser.add_argument("--visdom_server", type=str, default="http://localhost")
  35. parser.add_argument("--no_visdom", action="store_true", help= \
  36. "Disable visdom.")
  37. args = parser.parse_args()
  38. # Run the training
  39. print_args(args, parser)
  40. train(**vars(args))