def test_fastspeech_gpu_trainable_and_decodable(model_dict):
# make args
idim, odim = 10, 25
teacher_model_args = make_transformer_args(**model_dict)
model_args = make_feedforward_transformer_args(**model_dict)
# setup batch
ilens = [10, 5]
olens = [20, 15]
device = torch.device('cuda')
batch = prepare_inputs(idim,
odim,
ilens,
olens,
model_args["spk_embed_dim"],
device=device)
# define teacher model and save it
teacher_model = Transformer(idim, odim, Namespace(**teacher_model_args))
tmpdir = tempfile.mkdtemp(prefix="tmp_", dir="/tmp")
torch.save(teacher_model.state_dict(), tmpdir + "/model.dummy.best")
with open(tmpdir + "/model.json", 'wb') as f:
f.write(
json.dumps((idim, odim, teacher_model_args),
indent=4,
ensure_ascii=False,
sort_keys=True).encode('utf_8'))
# define model
model_args["teacher_model"] = tmpdir + "/model.dummy.best"
model = FeedForwardTransformer(idim, odim, Namespace(**model_args))
model.to(device)
optimizer = torch.optim.Adam(model.parameters())
# trainable
loss = model(**batch).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
# decodable
model.eval()
with torch.no_grad():
if model_args["spk_embed_dim"] is None:
spemb = None
else:
spemb = batch["spembs"][0]
model.inference(batch["xs"][0][:batch["ilens"][0]], None, spemb=spemb)
model.calculate_all_attentions(**batch)
# remove tmpdir
if os.path.exists(tmpdir):
shutil.rmtree(tmpdir)
def test_fastspeech_inference(alpha):
# make args
idim, odim = 10, 25
model_args = make_feedforward_transformer_args()
# setup batch
ilens = [10, 5]
olens = [20, 15]
batch = prepare_inputs(idim, odim, ilens, olens, model_args["spk_embed_dim"])
# define model
model = FeedForwardTransformer(idim, odim, Namespace(**model_args))
# test inference
inference_args = Namespace(**{"fastspeech_alpha": alpha})
model.eval()
with torch.no_grad():
if model_args["spk_embed_dim"] is None:
spemb = None
else:
spemb = batch["spembs"][0]
model.inference(
batch["xs"][0][: batch["ilens"][0]], inference_args, spemb=spemb,
)