def load_model(ckpt_pth): ckpt_dict = torch.load(ckpt_pth) model = Tacotron2() model.load_state_dict(ckpt_dict['model']) for name, param in model.named_parameters(): print(name) model = mode(model, True).eval() return model
def load_model(ckpt_pth): ckpt_dict = torch.load(ckpt_pth) model = Tacotron2() model.load_state_dict(ckpt_dict['model']) model = mode(model, True).eval() model.decoder.train() model.postnet.train() return model
def run(args): """Runs the algorithm.""" Path(hp.output_path).mkdir(parents=True, exist_ok=True) # setup nnabla context ctx = get_extension_context(args.context, device_id='0') nn.set_default_context(ctx) hp.comm = CommunicatorWrapper(ctx) hp.event = StreamEventHandler(int(hp.comm.ctx.device_id)) if hp.comm.n_procs > 1 and hp.comm.rank == 0: n_procs = hp.comm.n_procs logger.info(f'Distributed training with {n_procs} processes.') rng = np.random.RandomState(hp.seed) # train data train_loader = data_iterator(LJSpeechDataSource('metadata_train.csv', hp, shuffle=True, rng=rng), batch_size=hp.batch_size, with_memory_cache=False) # valid data valid_loader = data_iterator(LJSpeechDataSource('metadata_valid.csv', hp, shuffle=False, rng=rng), batch_size=hp.batch_size, with_memory_cache=False) dataloader = dict(train=train_loader, valid=valid_loader) model = Tacotron2(hp) # setup optimizer anneal_steps = [ x * (train_loader.size // hp.batch_size) for x in hp.anneal_steps ] lr_scheduler = AnnealingScheduler(hp.alpha, warmup=hp.warmup, anneal_steps=anneal_steps, anneal_factor=hp.anneal_factor) optimizer = Optimizer(weight_decay=hp.weight_decay, max_norm=hp.max_norm, lr_scheduler=lr_scheduler, name='Adam', alpha=hp.alpha) Tacotron2Trainer(model, dataloader, optimizer, hp).run()
def synthesize(args): char2idx = {ch: i for i, ch in enumerate(hp.vocab)} with open(args.f_text, 'r') as file: text = ''.join(file.readlines()) # normalize the text text = text_normalize(text, hp.vocab) if len(text) >= hp.text_len - 1: text = text[:hp.text_len - 1] text += '~' * (hp.text_len - len(text)) text = np.array([char2idx[ch] for ch in text]).reshape(-1) hp.batch_size = 1 # load the model model = Tacotron2(hp) model.training = False model.load_parameters(args.f_model) x_txt = nn.Variable([hp.batch_size, hp.text_len]) _, mels, _, _ = model(x_txt) x_txt.d = text[np.newaxis, :] mels.forward(clear_buffer=True) m = mels.d.copy().reshape(1, -1, hp.n_mels) np.save(args.f_output, m.transpose((0, 2, 1)))
def train(args): # build model model = Tacotron2() mode(model, True) optimizer = torch.optim.Adam(model.parameters(), lr = hps.lr, betas = hps.betas, eps = hps.eps, weight_decay = hps.weight_decay) criterion = Tacotron2Loss() # load checkpoint iteration = 1 if args.ckpt_pth != '': model, optimizer, iteration = load_checkpoint(args.ckpt_pth, model, optimizer) iteration += 1 # next iteration is iteration+1 # get scheduler if hps.sch: lr_lambda = lambda step: hps.sch_step**0.5*min((step+1)*hps.sch_step**-1.5, (step+1)**-0.5) if args.ckpt_pth != '': scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch = iteration) else: scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) # make dataset train_loader = prepare_dataloaders(args.data_dir) # get logger ready if args.log_dir != '': if not os.path.isdir(args.log_dir): os.makedirs(args.log_dir) os.chmod(args.log_dir, 0o775) logger = Tacotron2Logger(args.log_dir) # get ckpt_dir ready if args.ckpt_dir != '' and not os.path.isdir(args.ckpt_dir): os.makedirs(args.ckpt_dir) os.chmod(args.ckpt_dir, 0o775) model.train() # ================ MAIN TRAINNIG LOOP! =================== while iteration <= hps.max_iter: for batch in train_loader: if iteration > hps.max_iter: break start = time.perf_counter() x, y = model.parse_batch(batch) y_pred = model(x) # loss loss, item = criterion(y_pred, y, iteration) # zero grad model.zero_grad() # backward, grad_norm, and update loss.backward() grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), hps.grad_clip_thresh) optimizer.step() if hps.sch: scheduler.step() # info dur = time.perf_counter()-start print('Iter: {} Loss: {:.2e} Grad Norm: {:.2e} {:.1f}s/it'.format( iteration, item, grad_norm, dur)) # log if args.log_dir != '' and (iteration % hps.iters_per_log == 0): learning_rate = optimizer.param_groups[0]['lr'] logger.log_training(item, grad_norm, learning_rate, iteration) # sample if args.log_dir != '' and (iteration % hps.iters_per_sample == 0): model.eval() output = infer(hps.eg_text, model) model.train() logger.sample_training(output, iteration) # save ckpt if args.ckpt_dir != '' and (iteration % hps.iters_per_ckpt == 0): ckpt_pth = os.path.join(args.ckpt_dir, 'ckpt_{}'.format(iteration)) save_checkpoint(model, optimizer, iteration, ckpt_pth) iteration += 1 if args.log_dir != '': logger.close()