def test_fastspeech_gpu_trainable_and_decodable(teacher_type, model_dict):
# make args
idim, odim = 10, 25
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
if teacher_type == "transformer":
teacher_model_args = make_transformer_args(**model_dict)
teacher_model = Transformer(idim, odim,
Namespace(**teacher_model_args))
elif teacher_type == "tacotron2":
teacher_model_args = make_taco2_args(**model_dict)
teacher_model = Tacotron2(idim, odim, Namespace(**teacher_model_args))
else:
raise ValueError()
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_multi_gpu_trainable(teacher_type, model_dict):
# make args
idim, odim = 10, 25
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
if teacher_type == "transformer":
teacher_model_args = make_transformer_args(**model_dict)
teacher_model = Transformer(idim, odim,
Namespace(**teacher_model_args))
elif teacher_type == "tacotron2":
teacher_model_args = make_taco2_args(**model_dict)
teacher_model = Tacotron2(idim, odim, Namespace(**teacher_model_args))
else:
raise ValueError()
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
ngpu = 2
device_ids = list(range(ngpu))
model_args["teacher_model"] = tmpdir + "/model.dummy.best"
model = FeedForwardTransformer(idim, odim, Namespace(**model_args))
model = torch.nn.DataParallel(model, device_ids)
model.to(device)
optimizer = torch.optim.Adam(model.parameters())
# trainable
loss = model(**batch).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
# remove tmpdir
if os.path.exists(tmpdir):
shutil.rmtree(tmpdir)