def test_tacotron2_multi_gpu_trainable(model_dict):
ngpu = 2
device_ids = list(range(ngpu))
device = torch.device('cuda')
bs = 10
maxin_len = 10
maxout_len = 10
idim = 5
odim = 10
model_args = make_taco2_args(**model_dict)
batch = prepare_inputs(bs, idim, odim, maxin_len, maxout_len,
model_args['spk_embed_dim'], model_args['spc_dim'],
device=device)
# define model
model = Tacotron2(idim, odim, Namespace(**model_args))
model = torch.nn.DataParallel(model, device_ids)
optimizer = torch.optim.Adam(model.parameters())
model.to(device)
# trainable
loss = model(**batch).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
def test_tacotron2_trainable_and_decodable(model_dict, inference_dict):
# make args
model_args = make_taco2_args(**model_dict)
inference_args = make_inference_args(**inference_dict)
# setup batch
bs = 2
maxin_len = 10
maxout_len = 10
idim = 5
odim = 10
if model_args['use_cbhg']:
model_args['spc_dim'] = 129
if model_args['use_speaker_embedding']:
model_args['spk_embed_dim'] = 128
batch = prepare_inputs(bs, idim, odim, maxin_len, maxout_len,
model_args['spk_embed_dim'], model_args['spc_dim'])
# define model
model = Tacotron2(idim, odim, Namespace(**model_args))
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():
spemb = None if model_args['spk_embed_dim'] is None else batch["spembs"][0]
model.inference(batch["xs"][0][:batch["ilens"][0]], Namespace(**inference_args), spemb)
model.calculate_all_attentions(**batch)
def test_tacotron2_gpu_trainable_and_decodable(model_dict):
bs = 2
maxin_len = 10
maxout_len = 10
idim = 5
odim = 10
device = torch.device('cuda')
model_args = make_taco2_args(**model_dict)
inference_args = make_inference_args()
batch = prepare_inputs(bs, idim, odim, maxin_len, maxout_len,
model_args['spk_embed_dim'], model_args['spc_dim'],
device=device)
# define model
model = Tacotron2(idim, odim, Namespace(**model_args))
optimizer = torch.optim.Adam(model.parameters())
model.to(device)
# trainable
loss = model(**batch).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
# decodable
model.eval()
with torch.no_grad():
spemb = None if model_args['spk_embed_dim'] is None else batch["spembs"][0]
model.inference(batch["xs"][0][:batch["ilens"][0]], Namespace(**inference_args), spemb)
model.calculate_all_attentions(**batch)
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)
