Beginning

train.py

@@ -0,0 +1,180 @@
+import os
+import sys
+import time
+import torch
+import signal
+import argparse
+import numpy as np
+
+import torch.nn as nn
+from torch import optim
+from torch.utils.data import DataLoader
+
+from network import *
+import train_config as c
+from utils.generic_utils import (Progbar, remove_experiment_folder,
+                                 create_experiment_folder, save_checkpoint)
+from utils.model import get_param_size
+from datasets.LJSpeech import LJSpeechDataset
+
+use_cuda = torch.cuda.is_available()
+
+_ = os.path.dirname(os.path.realpath(__file__))
+OUT_PATH = os.path.join(_, c.output_path)
+OUT_PATH = create_experiment_folder(OUT_PATH)
+
+def signal_handler(signal, frame):
+    print(" !! Pressed Ctrl+C !!")
+    remove_experiment_folder(OUT_PATH)
+    sys.exit(0)
+
+
+def main(args):
+
+    dataset = LJSpeechDataset(os.path.join(c.data_path, 'metadata.csv'),
+                              os.path.join(c.data_path, 'wavs'),
+                              c.dec_out_per_step
+                             )
+
+    model = Tacotron(c.embedding_size,
+                     c.hidden_size,
+                     c.num_mels,
+                     c.num_freq,
+                     c.dec_out_per_step,
+                     c.teacher_forcing_ratio)
+    if use_cuda:
+        model = nn.DataParallel(model.cuda())
+
+    optimizer = optim.Adam(model.parameters(), lr=c.lr)
+
+    try:
+        checkpoint = torch.load(os.path.join(
+            c.checkpoint_path, 'checkpoint_%d.pth.tar' % args.restore_step))
+        model.load_state_dict(checkpoint['model'])
+        optimizer.load_state_dict(checkpoint['optimizer'])
+        print("\n > Model restored from step %d\n" % args.restore_step)
+
+    except:
+        print("\n > Starting a new training\n")
+
+    model = model.train()
+
+    if not os.path.exists(c.checkpoint_path):
+        os.mkdir(c.checkpoint_path)
+
+    if use_cuda:
+        criterion = nn.L1Loss().cuda()
+    else:
+        criterion = nn.L1Loss()
+
+    n_priority_freq = int(3000 / (c.sample_rate * 0.5) * c.num_freq)
+
+    for epoch in range(c.epochs):
+
+        dataloader = DataLoader(dataset, batch_size=args.batch_size,
+                                shuffle=True, collate_fn=dataset.collate_fn,
+                                drop_last=True, num_workers=8)
+        progbar = Progbar(len(dataset) / args.batch_size)
+
+        for i, data in enumerate(dataloader):
+
+            current_step = i + args.restore_step + epoch * len(dataloader) + 1
+
+            optimizer.zero_grad()
+
+            try:
+                mel_input = np.concatenate((np.zeros(
+                    [args.batch_size, c.num_mels, 1], dtype=np.float32), data[2][:, :, 1:]), axis=2)
+            except:
+                raise TypeError("not same dimension")
+
+            if use_cuda:
+                characters = Variable(torch.from_numpy(data[0]).type(
+                    torch.cuda.LongTensor), requires_grad=False).cuda()
+                mel_input = Variable(torch.from_numpy(mel_input).type(
+                    torch.cuda.FloatTensor), requires_grad=False).cuda()
+                mel_spectrogram = Variable(torch.from_numpy(data[2]).type(
+                    torch.cuda.FloatTensor), requires_grad=False).cuda()
+                linear_spectrogram = Variable(torch.from_numpy(data[1]).type(
+                    torch.cuda.FloatTensor), requires_grad=False).cuda()
+
+            else:
+                characters = Variable(torch.from_numpy(data[0]).type(
+                    torch.LongTensor), requires_grad=False)
+                mel_input = Variable(torch.from_numpy(mel_input).type(
+                    torch.FloatTensor), requires_grad=False)
+                mel_spectrogram = Variable(torch.from_numpy(
+                    data[2]).type(torch.FloatTensor), requires_grad=False)
+                linear_spectrogram = Variable(torch.from_numpy(
+                    data[1]).type(torch.FloatTensor), requires_grad=False)
+
+            mel_output, linear_output = model.forward(characters, mel_input)
+
+            mel_loss = criterion(mel_output, mel_spectrogram)
+            linear_loss = torch.abs(linear_output - linear_spectrogram)
+            linear_loss = 0.5 * \
+                torch.mean(linear_loss) + 0.5 * \
+                torch.mean(linear_loss[:, :n_priority_freq, :])
+            loss = mel_loss + linear_loss
+            loss = loss.cuda()
+
+            start_time = time.time()
+
+            loss.backward()
+
+            nn.utils.clip_grad_norm(model.parameters(), 1.)
+
+            optimizer.step()
+
+            time_per_step = time.time() - start_time
+            progbar.update(i, values=[('total_loss', loss.data[0]),
+                                      ('linear_loss', linear_loss.data[0]),
+                                      ('mel_loss', mel_loss.data[0])])
+
+            if current_step % c.save_step == 0:
+                checkpoint_path = 'checkpoint_{}.pth.tar'.format(current_step)
+                checkpoint_path = os.path.join(OUT_PATH, checkpoint_path)
+                save_checkpoint({'model': model.state_dict(),
+                                 'optimizer': optimizer.state_dict(),
+                                 'step': current_step,
+                                 'total_loss': loss.data[0],
+                                 'linear_loss': linear_loss.data[0],
+                                 'mel_loss': mel_loss.data[0],
+                                 'date': datetime.date.today().strftime("%B %d, %Y")},
+                                checkpoint_path)
+                print(" > Checkpoint is saved : {}".format(checkpoint_path))
+
+            if current_step in c.decay_step:
+                optimizer = adjust_learning_rate(optimizer, current_step)
+
+
+def adjust_learning_rate(optimizer, step):
+    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
+    if step == 500000:
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = 0.0005
+
+    elif step == 1000000:
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = 0.0003
+
+    elif step == 2000000:
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = 0.0001
+
+    return optimizer
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--restore_step', type=int,
+                        help='Global step to restore checkpoint', default=0)
+    parser.add_argument('--batch_size', type=int,
+                        help='Batch size', default=32)
+    parser.add_argument('--config', type=str,
+                       help='path to config file for training',)
+    args = parser.parse_args()
+
+    signal.signal(signal.SIGINT, signal_handler)
+
+    main(args)