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()