Batch update after data-loss

train.py

@@ -14,16 +14,16 @@ from torch.utils.data import DataLoader
 from tensorboardX import SummaryWriter
 
 from utils.generic_utils import (
-    synthesis, remove_experiment_folder, create_experiment_folder,
+    remove_experiment_folder, create_experiment_folder,
     save_checkpoint, save_best_model, load_config, lr_decay, count_parameters,
     check_update, get_commit_hash, sequence_mask, AnnealLR)
 from utils.visual import plot_alignment, plot_spectrogram
 from models.tacotron import Tacotron
 from layers.losses import L1LossMasked
 from utils.audio import AudioProcessor
+from utils.synthesis import synthesis
 
 torch.manual_seed(1)
-# torch.set_num_threads(4)
 use_cuda = torch.cuda.is_available()
 
 
@@ -36,7 +36,7 @@ def train(model, criterion, criterion_st, data_loader, optimizer, optimizer_st,
     avg_stop_loss = 0
     avg_step_time = 0
     print(" | > Epoch {}/{}".format(epoch, c.epochs), flush=True)
-    n_priority_freq = int(3000 / (c.sample_rate * 0.5) * c.num_freq)
+    n_priority_freq = int(3000 / (c.audio['sample_rate'] * 0.5) * c.audio['num_freq'])
     batch_n_iter = int(len(data_loader.dataset) / c.batch_size)
     for num_iter, data in enumerate(data_loader):
         start_time = time.time()
@@ -215,7 +215,7 @@ def evaluate(model, criterion, criterion_st, data_loader, ap, current_step):
         "I'm sorry Dave. I'm afraid I can't do that.",
         "This cake is great. It's so delicious and moist."
     ]
-    n_priority_freq = int(3000 / (c.sample_rate * 0.5) * c.num_freq)
+    n_priority_freq = int(3000 / (c.audio['sample_rate'] * 0.5) * c.audio['num_freq'])
     with torch.no_grad():
         if data_loader is not None:
             for num_iter, data in enumerate(data_loader):
@@ -277,6 +277,7 @@ def evaluate(model, criterion, criterion_st, data_loader, ap, current_step):
             const_spec = linear_output[idx].data.cpu().numpy()
             gt_spec = linear_input[idx].data.cpu().numpy()
             align_img = alignments[idx].data.cpu().numpy()
+
             const_spec = plot_spectrogram(const_spec, ap)
             gt_spec = plot_spectrogram(gt_spec, ap)
             align_img = plot_alignment(align_img)
@@ -319,8 +320,8 @@ def evaluate(model, criterion, criterion_st, data_loader, ap, current_step):
     ap.griffin_lim_iters = 60
     for idx, test_sentence in enumerate(test_sentences):
         try:
-            wav, linear_spec, alignments = synthesis(model, ap, test_sentence,
-                                                     use_cuda, c.text_cleaner)
+            wav, alignment, linear_spec, stop_tokens = synthesis(model, test_sentence, c,
+                                                     use_cuda, ap)
 
             file_path = os.path.join(AUDIO_PATH, str(current_step))
             os.makedirs(file_path, exist_ok=True)
@@ -330,7 +331,7 @@ def evaluate(model, criterion, criterion_st, data_loader, ap, current_step):
 
             wav_name = 'TestSentences/{}'.format(idx)
             tb.add_audio(
-                wav_name, wav, current_step, sample_rate=c.sample_rate)
+                wav_name, wav, current_step, sample_rate=c.audio['sample_rate'])
             align_img = alignments[0].data.cpu().numpy()
             linear_spec = plot_spectrogram(linear_spec, ap)
             align_img = plot_alignment(align_img)
@@ -345,28 +346,22 @@ def evaluate(model, criterion, criterion_st, data_loader, ap, current_step):
 
 
 def main(args):
-    dataset = importlib.import_module('datasets.' + c.dataset)
-    Dataset = getattr(dataset, 'MyDataset')
-    audio = importlib.import_module('utils.' + c.audio_processor)
+    preprocessor = importlib.import_module('datasets.preprocess')
+    preprocessor = getattr(preprocessor, c.dataset.lower())
+    MyDataset = importlib.import_module('datasets.'+c.data_loader)
+    MyDataset = getattr(MyDataset, "MyDataset")
+    audio = importlib.import_module('utils.' + c.audio['audio_processor'])
     AudioProcessor = getattr(audio, 'AudioProcessor')
 
-    ap = AudioProcessor(
-        sample_rate=c.sample_rate,
-        num_mels=c.num_mels,
-        min_level_db=c.min_level_db,
-        frame_shift_ms=c.frame_shift_ms,
-        frame_length_ms=c.frame_length_ms,
-        ref_level_db=c.ref_level_db,
-        num_freq=c.num_freq,
-        power=c.power,
-        preemphasis=c.preemphasis)
+    ap = AudioProcessor(**c.audio)
 
     # Setup the dataset
-    train_dataset = Dataset(
+    train_dataset = MyDataset(
         c.data_path,
         c.meta_file_train,
         c.r,
         c.text_cleaner,
+        preprocessor=preprocessor,
         ap=ap,
         batch_group_size=8*c.batch_size,
         min_seq_len=c.min_seq_len)
@@ -381,8 +376,8 @@ def main(args):
         pin_memory=True)
 
     if c.run_eval:
-        val_dataset = Dataset(
-            c.data_path, c.meta_file_val, c.r, c.text_cleaner, ap=ap, batch_group_size=0)
+        val_dataset = MyDataset(
+            c.data_path, c.meta_file_val, c.r, c.text_cleaner, preprocessor=preprocessor, ap=ap, batch_group_size=0)
 
         val_loader = DataLoader(
             val_dataset,
@@ -395,7 +390,7 @@ def main(args):
     else:
         val_loader = None
 
-    model = Tacotron(c.embedding_size, ap.num_freq, c.num_mels, c.r)
+    model = Tacotron(c.embedding_size, ap.num_freq, ap.num_mels, c.r)
     print(" | > Num output units : {}".format(ap.num_freq), flush=True)
 
     optimizer = optim.Adam(model.parameters(), lr=c.lr, weight_decay=0)
@@ -431,7 +426,7 @@ def main(args):
             criterion.cuda()
             criterion_st.cuda()
 
-    scheduler = AnnealLR(optimizer, warmup_steps=c.warmup_steps, last_epoch= (args.restore_step-1))
+    scheduler = AnnealLR(optimizer, warmup_steps=c.warmup_steps, last_epoch=args.restore_step - 1)
     num_params = count_parameters(model)
     print(" | > Model has {} parameters".format(num_params), flush=True)
 
@@ -454,7 +449,7 @@ def main(args):
         best_loss = save_best_model(model, optimizer, train_loss, best_loss,
                                     OUT_PATH, current_step, epoch)
          # shuffle batch groups
-        train_loader.dataset.sort_frames()
+        train_loader.dataset.sort_items()
 
 
 if __name__ == '__main__':
@@ -477,8 +472,9 @@ if __name__ == '__main__':
     parser.add_argument(
         '--data_path',
         type=str,
-        default='',
-        help='data path to overrite config.json')
+        help='dataset path.',
+        default=''
+    )
     args = parser.parse_args()
 
     # setup output paths and read configs
@@ -491,7 +487,7 @@ if __name__ == '__main__':
     os.makedirs(AUDIO_PATH, exist_ok=True)
     shutil.copyfile(args.config_path, os.path.join(OUT_PATH, 'config.json'))
 
-    if args.data_path != "":
+    if args.data_path != '':
         c.data_path = args.data_path
 
     # setup tensorboard