mtts/train.py

import argparse
import os

import numpy as np
import torch
import yaml
from torch.optim import SGD, Adam
from torch.optim.lr_scheduler import CyclicLR, OneCycleLR, ReduceLROnPlateau
from torch.utils.data import BatchSampler, DataLoader, SubsetRandomSampler
from torch.utils.tensorboard import SummaryWriter

from mtts.datasets.dataset import Dataset, collate_fn
from mtts.loss import FS2Loss
from mtts.models.fs2_model import FastSpeech2
from mtts.optimizer import ScheduledOptim
from mtts.utils.logging import get_logger
from mtts.utils.utils import save_image

logger = get_logger(__file__)


class AverageMeter:
    def __init__(self):
        self.mel_loss_v = 0.0
        self.posnet_loss_v = 0.0
        self.d_loss_v = 0.0
        self.total_loss_v = 0.0

        self._i = 0

    def update(self, mel_loss, posnet_loss, d_loss, total_loss):
        self.mel_loss_v = ((self.mel_loss_v * self._i) + mel_loss.item()) / (self._i + 1)
        self.posnet_loss_v = ((self.posnet_loss_v * self._i) + posnet_loss.item()) / (self._i + 1)
        self.d_loss_v = ((self.d_loss_v * self._i) + d_loss.item()) / (self._i + 1)
        self.total_loss_v = ((self.total_loss_v * self._i) + total_loss.item()) / (self._i + 1)

        self._i += 1
        return self.mel_loss_v, self.posnet_loss_v, self.d_loss_v, self.total_loss_v


def split_batch(data, i, n_split):
    n = data[1].shape[0]
    k = n // n_split
    ds = [d[:, i * k:(i + 1) * k] if j == 0 else d[i * k:(i + 1) * k] for j, d in enumerate(data)]
    return ds


def shuffle(data):
    n = data[1].shape[0]
    idx = np.random.permutation(n)
    data_shuffled = [d[:, idx] if i == 0 else d[idx] for i, d in enumerate(data)]
    return data_shuffled


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('-r', '--restore', type=str, default='')
    parser.add_argument('-c', '--config', type=str, default='./config.yaml')
    parser.add_argument('-d', '--device', type=str, default='cuda')
    args = parser.parse_args()

    device = args.device
    logger.info(f'using device {device}')

    with open(args.config) as f:
        config = yaml.safe_load(f)
        logger.info(f.read())

    dataset = Dataset(config)

    dataloader = DataLoader(dataset,
                            batch_size=config['training']['batch_size'],
                            shuffle=False,
                            collate_fn=collate_fn,
                            drop_last=False,
                            num_workers=config['training']['num_workers'])

    step_per_epoch = len(dataloader) * config['training']['batch_size']

    model = FastSpeech2(config)
    model = model.to(args.device)
    #model.encoder.emb_layers.to(device) # ?

    optim_conf = config['optimizer']
    optim_class = eval(optim_conf['type'])
    logger.info(optim_conf['params'])
    optimizer = optim_class(model.parameters(), **optim_conf['params'])

    if args.restore != '':
        logger.info(f'Loading checkpoint {args.restore}')
        content = torch.load(args.restore)
        model.load_state_dict(content['model'])
        optimizer.load_state_dict(content['optimizer'])
        current_step = content['step']
        start_epoch = current_step // step_per_epoch
        logger.info(f'loaded checkpoint at step {current_step}, epoch {start_epoch}')
    else:
        current_step = 0
        start_epoch = 0
        logger.info(f'Start training from scratch,step={current_step},epoch={start_epoch}')

    lrs = np.linspace(0, optim_conf['params']['lr'], optim_conf['n_warm_up_step'])
    Scheduler = eval(config['lr_scheduler']['type'])
    lr_scheduler = Scheduler(optimizer, **config['lr_scheduler']['params'])

    loss_fn = FS2Loss().to(device)
    train_logger = SummaryWriter(config['training']['log_path'])
    val_logger = SummaryWriter(config['training']['log_path'])
    avg = AverageMeter()
    for epoch in range(start_epoch, config['training']['epochs']):

        model.train()
        for i, data in enumerate(dataloader):
            data = shuffle(data)
            max_src_len = torch.max(data[-2])
            max_mel_len = torch.max(data[-1])
            for k in range(config['training']['batch_split']):
                data_split = split_batch(data, k, config['training']['batch_split'])
                tokens, duration, mel_truth, seq_len, mel_len = data_split
                #print(mel_len)
                tokens = tokens.to(device)
                duration = duration.to(device)
                mel_truth = mel_truth.to(device)
                seq_len = seq_len.to(device)
                mel_len = mel_len.to(device)
                # if torch.max(log_D) > 50:
                #  logger.info('skipping sample')
                #  continue

                mel_truth = mel_truth - config['fbank']['mel_mean']
                duration = duration - config['duration_predictor']['duration_mean']
                output = model(tokens, seq_len, mel_len, duration, max_src_len=max_src_len, max_mel_len=max_mel_len)

                mel_pred, mel_postnet, d_pred, src_mask, mel_mask, mel_len = output

                mel_loss, mel_postnet_loss, d_loss = loss_fn(d_pred, duration, mel_pred, mel_postnet, mel_truth,
                                                             ~src_mask, ~mel_mask)

                total_loss = mel_postnet_loss + d_loss + mel_loss
                ml, pl, dl, tl = avg.update(mel_loss, mel_postnet_loss, d_loss, total_loss)
                lr = optimizer.param_groups[0]['lr']
                msg = f'epoch:{epoch},step:{current_step}|{step_per_epoch},loss:{tl:.3},mel:{ml:.3},'
                msg += f'mel_postnet:{pl:.3},duration:{dl:.3},{lr:.3}'

                if current_step % config['training']['log_step'] == 0:
                    logger.info(msg)

                total_loss = total_loss / config['training']['acc_step']
                total_loss.backward()
                if current_step % config['training']['acc_step'] != 0:
                    continue

                current_step += 1
                if current_step < config['optimizer']['n_warm_up_step']:
                    lr = lrs[current_step]
                    optimizer.param_groups[0]['lr'] = lr
                    optimizer.step()
                    optimizer.zero_grad()
                else:
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()

                if current_step % config['training']['synth_step'] == 0:
                    mel_pred = mel_pred.detach().cpu().numpy()
                    mel_truth = mel_truth.detach().cpu().numpy()
                    saved_path = os.path.join(config['training']['log_path'], f'{current_step}.png')
                    save_image(mel_truth[0][:mel_len[0]], mel_pred[0][:mel_len[0]], saved_path)
                    np.save(saved_path + '.npy', mel_pred[0])

                if current_step % config['training']['log_step'] == 0:

                    train_logger.add_scalar('total_loss', tl, current_step)
                    train_logger.add_scalar('mel_loss', ml, current_step)
                    train_logger.add_scalar('mel_postnet_loss', pl, current_step)
                    train_logger.add_scalar('duration_loss', dl, current_step)

                if current_step % config['training']['checkpoint_step'] == 0:
                    if not os.path.exists(config['training']['checkpoint_path']):
                        os.makedirs(config['training']['checkpoint_path'])
                    ckpt_file = os.path.join(config['training']['checkpoint_path'],
                                             'checkpoint_{}.pth.tar'.format(current_step))
                    content = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'step': current_step}

                    torch.save(content, ckpt_file)
                    logger.info(f'Saved model at step {current_step} to {ckpt_file}')
    logger.info(f"End of training for epoch {config['training']['epochs']}")