add train, check notebook

neo-tts/tacotron2-cac/README.md

@@ -72,6 +72,7 @@ Train your model with
 ```
 python3 train.py --dataset DATASET
 ```
+Check `training.ipynb` or the up-to-date [Colab Notebook](https://colab.research.google.com/drive/1ra5DRIyohXSNiT-3miBWkTZbaGDTczOJ?usp=sharing).
 
 # TensorBoard
 
@@ -86,7 +87,12 @@ Use
 python3 infer_arpa.py --mode single --dataset SingleSpk22 --checkpoint tacotron2.pth --out-path MyAud --text "there are no faces"
 ```
 Or check `inference.ipynb` for more lean, hands-on notebook version
-
+# Checking model
+Use
+```
+python3 check_model.py --dataset SingleSpk22 --checkpoint tacotron2.pth
+```
+This script will try to load a checkpoint to CPU. If the file doesn't exist, it will exit with `2` return code. If it does but there's an error loading, it will exit with `1`. Otherwise, if the model is loaded well, it will exit with `0`
 
 to serve TensorBoard on your localhost.
 <!-- The loss curves, synthesized mel-spectrograms, and audios are shown.

neo-tts/tacotron2-cac/check_model.py

@@ -0,0 +1,73 @@
+import os
+import json
+# import re
+import argparse
+# from string import punctuation
+
+import torch
+import numpy as np
+from torch.utils.data import DataLoader
+# from g2p_en import G2p
+# from pypinyin import pinyin, Style
+
+from utils.model import get_model, get_vocoder
+from utils.tools import get_configs_of, to_device #, read_lexicon
+from dataset import TextDataset
+from text import text_to_sequence, sequence_to_text
+from text.cleaners import to_arpa
+import sys
+device = "cpu"
+from model import Tacotron2
+
+def get_tac2(configs, device):
+    (preprocess_config, model_config, train_config) = configs
+
+    model = Tacotron2(preprocess_config, model_config, train_config).to(device)
+
+    model.eval()
+    model.requires_grad_ = False
+    return model
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        required=True,
+        help="name of dataset",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        required=True,
+        help="Path to model file.",
+    )
+    args = parser.parse_args()
+
+    if not os.path.isfile(args.checkpoint):
+      print(f"Model path {args.checkpoint} incorrect, file not found.")
+      sys.exit(2)
+
+
+
+    # Read Config
+    preprocess_config, model_config, train_config = get_configs_of(args.dataset)
+    configs = (preprocess_config, model_config, train_config)
+    stats = None
+    mel_stats = None
+
+    # Get model
+    model = get_tac2(configs, device)
+    try:
+        ckpt = torch.load(args.checkpoint)
+        model.load_state_dict(ckpt["model"])
+    except:
+        print("Error loading model")
+        sys.exit(1) # Error
+    
+    print("Model is valid.")
+    sys.exit(0) # Success
+    
+ 

neo-tts/tacotron2-cac/train.py

@@ -48,7 +48,11 @@ def main(args, configs):
         print(f"Warm starting from checkpoint {args.pretrained}")
         ckpt = torch.load(args.pretrained)
         model.load_state_dict(ckpt["model"])
-        
+    
+    if args.checkpoint_interval > 1:
+      print(f"Overriding checkpoint interval to {args.checkpoint_interval}")
+      train_config["step"]["save_step"] = args.checkpoint_interval
+
     model = nn.DataParallel(model)
     num_param = get_param_num(model)
     Loss = Tacotron2Loss(preprocess_config, model_config, train_config).to(device)
@@ -60,6 +64,12 @@ def main(args, configs):
     vocoder = get_vocoder(model_config, device)
 
     # Init logger
+    
+    if len(args.out_path):
+        train_config["path"]["log_path"] = os.path.join(args.out_path, "logs")
+        train_config["path"]["ckpt_path"] = os.path.join(args.out_path, "checkpoints")
+        train_config["path"]["result_path"] = os.path.join(args.out_path, "val")
+        
     for p in train_config["path"].values():
         os.makedirs(p, exist_ok=True)
     train_log_path = os.path.join(train_config["path"]["log_path"], "train")
@@ -211,6 +221,20 @@ if __name__ == "__main__":
         help="Path to pretrained checkpoint file to finetune from",
         default="",
     )
+    parser.add_argument(
+        "--out-path",
+        type=str,
+        required=False,
+        help="Path override for train",
+        default="",
+    )
+    parser.add_argument(
+        "--checkpoint-interval",
+        type=int,
+        required=False,
+        help="Checkpoint interval override for train",
+        default=0,
+    )
     args = parser.parse_args()
 
     # Read Config