def main(argv):
args = utils.parse_args("Train a wavenet model")
utils.redirect_log_to_file(args.model_dir)
hparams = create_hparams(args.model_dir, args.configs, initialize=True)
utils.check_git_hash(args.model_dir)
# Prepare data
data.load_vocab(hparams)
train_input_fn = data.InputPipeline(hparams, tf.estimator.ModeKeys.TRAIN)
eval_input_fn = data.InputPipeline(hparams, tf.estimator.ModeKeys.EVAL)
# Training
train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn,
max_steps=hparams.train_steps)
eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn,
steps=hparams.eval_steps,
throttle_secs=hparams.throttle_secs)
distribution = tf.contrib.distribute.MirroredStrategy()
run_config = tf.estimator.RunConfig(
model_dir=args.model_dir,
train_distribute=distribution,
save_summary_steps=hparams.save_summary_steps,
save_checkpoints_secs=hparams.save_checkpoints_secs,
keep_checkpoint_max=hparams.n_checkpoints)
estimator = tf.estimator.Estimator(model_fn=model.build_model_fn(hparams),
config=run_config,
model_dir=args.model_dir)
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
def main():
parser = utils.get_argument_parser("Decode by using the trained model")
parser.add_argument("--checkpoint", dest="checkpoint", help="Path to a checkpoint file. Default is the latest version.")
parser.add_argument("--limit", type=int, default=0, help="The number of sentences to be decoded. (default=unlimited)")
parser.add_argument("--use_eval", action="store_true", help="Use evaluation dataset for prediction")
parser.add_argument("--predict_file", type=str, default="", help="Path to a text file to be translated")
parser.add_argument("--out_file", type=str, default="", help="Path to a text file to write")
args = parser.parse_args()
hparams = create_hparams(args.model_dir, args.configs, initialize=False)
utils.check_git_hash(args.model_dir)
data.load_vocab(hparams)
if args.use_eval:
pipeline = data.InputPipeline(hparams.source_eval_file, None, None, tf.estimator.ModeKeys.PREDICT, hparams)
else:
pipeline = data.InputPipeline(args.predict_file, None, None, tf.estimator.ModeKeys.PREDICT, hparams)
estimator = tf.estimator.Estimator(model_fn=model.build_model_fn(hparams), model_dir=args.model_dir)
# set a file path to write
if args.out_file != "":
f = open(args.out_file, 'w')
else:
f = None
for i, prediction in enumerate(estimator.predict(pipeline, checkpoint_path=args.checkpoint)):
if args.limit and i == args.limit:
break
token_ids = prediction.tolist()
print(hparams.target_vocab.decode(token_ids), file=f)
if i % 1000 == 0:
tf.logging.info("write: %d", i)
def train_and_eval(rank, n_gpus, hps):
global global_step
if rank == 0:
logger = utils.get_logger(hps.model_dir)
logger.info(hps)
utils.check_git_hash(hps.model_dir)
writer = SummaryWriter(log_dir=hps.model_dir)
writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)
torch.manual_seed(hps.train.seed)
torch.cuda.set_device(rank)
train_dataset = TextMelLoader(hps.data.training_files, hps.data)
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset,
num_replicas=n_gpus,
rank=rank,
shuffle=True)
collate_fn = TextMelCollate(1)
train_loader = DataLoader(train_dataset, num_workers=8, shuffle=False,
batch_size=hps.train.batch_size, pin_memory=True,
drop_last=True, collate_fn=collate_fn, sampler=train_sampler)
if rank == 0:
val_dataset = TextMelLoader(hps.data.validation_files, hps.data)
val_loader = DataLoader(val_dataset, num_workers=8, shuffle=False,
batch_size=hps.train.batch_size, pin_memory=True,
drop_last=True, collate_fn=collate_fn)
generator = models.FlowGenerator(
n_vocab=len(symbols),
out_channels=hps.data.n_mel_channels,
**hps.model).cuda(rank)
optimizer_g = commons.Adam(generator.parameters(), scheduler=hps.train.scheduler, dim_model=hps.model.hidden_channels, warmup_steps=hps.train.warmup_steps, lr=hps.train.learning_rate, betas=hps.train.betas, eps=hps.train.eps)
if hps.train.fp16_run:
generator, optimizer_g._optim = amp.initialize(generator, optimizer_g._optim, opt_level="O1")
generator = DDP(generator)
epoch_str = 1
global_step = 0
try:
_, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), generator, optimizer_g)
epoch_str += 1
optimizer_g.step_num = (epoch_str - 1) * len(train_loader)
optimizer_g._update_learning_rate()
global_step = (epoch_str - 1) * len(train_loader)
except:
if hps.train.ddi and os.path.isfile(os.path.join(hps.model_dir, "ddi_G.pth")):
_ = utils.load_checkpoint(os.path.join(hps.model_dir, "ddi_G.pth"), generator, optimizer_g)
for epoch in range(epoch_str, hps.train.epochs + 1):
if rank==0:
train(rank, epoch, hps, generator, optimizer_g, train_loader, logger, writer)
evaluate(rank, epoch, hps, generator, optimizer_g, val_loader, logger, writer_eval)
if epoch%50 == 0:
utils.save_checkpoint(generator, optimizer_g, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "G_{}.pth".format(epoch)))
else:
train(rank, epoch, hps, generator, optimizer_g, train_loader, None, None)
def main():
parser = utils.get_argument_parser("Decode by using the trained model")
parser.add_argument(
"--checkpoint",
dest="checkpoint",
help="Path to a checkpoint file. Default is the latest version.")
parser.add_argument(
"--limit",
type=int,
default=1,
help="The number of sentences to be decoded. (0=unlimited)")
parser.add_argument("--use_eval",
action="store_true",
help="Use evaluation dataset for prediction")
parser.add_argument("--predict_dir",
type=str,
default="",
help="Path to a local condition file dir")
parser.add_argument("--out_dir",
type=str,
default="",
help="Path to a wav file dir to write")
args = parser.parse_args()
hparams = create_hparams(args.model_dir, args.configs, initialize=False)
utils.check_git_hash(args.model_dir)
if not os.path.isdir(args.out_dir):
os.mkdir(os.path.relpath(args.out_dir))
data.load_vocab(hparams)
if args.use_eval:
input_fn = data.InputPipeline(hparams, tf.estimator.ModeKeys.EVAL)
else:
raise NotImplementedError("File to mel or wav is not avaliable now.")
estimator = tf.estimator(model_fn=model.build_model_fn(hparams),
model_dir=args.model_dir)
for i, prediction in enumerate(
estimator.predict(input_fn, checkpoint_path=args.checkpoint)):
for j, wav in predictions.tolist():
wav = wav.astype(np.float32)
write(
os.path.join(args.out_dir,
"{}.wav".format(i * hparams.infer_batch_size +
i)), hparams.sample_rate, wav)
if args.limit and i + 1 == args.limit:
break
def main():
hps = utils.get_hparams()
logger = utils.get_logger(hps.model_dir)
logger.info(hps)
utils.check_git_hash(hps.model_dir)
torch.manual_seed(hps.train.seed)
train_dataset = TextMelLoader(hps.data.training_files, hps.data)
collate_fn = TextMelCollate(1)
train_loader = DataLoader(train_dataset,
num_workers=8,
shuffle=True,
batch_size=hps.train.batch_size,
pin_memory=True,
drop_last=True,
collate_fn=collate_fn)
generator = FlowGenerator_DDI(speaker_dim=hps.model.speaker_embedding,
n_vocab=len(symbols),
out_channels=hps.data.n_mel_channels,
**hps.model).cuda()
optimizer_g = commons.Adam(generator.parameters(),
scheduler=hps.train.scheduler,
dim_model=hps.model.hidden_channels,
warmup_steps=hps.train.warmup_steps,
lr=hps.train.learning_rate,
betas=hps.train.betas,
eps=hps.train.eps)
generator.train()
for batch_idx, (x, x_lengths, y, y_lengths,
speaker_embedding) in enumerate(train_loader):
x, x_lengths = x.cuda(), x_lengths.cuda()
y, y_lengths = y.cuda(), y_lengths.cuda()
speaker_embedding = speaker_embedding.cuda()
_ = generator(x, x_lengths, speaker_embedding, y, y_lengths, gen=False)
break
utils.save_checkpoint(generator, optimizer_g, hps.train.learning_rate, 0,
os.path.join(hps.model_dir, "ddi_G.pth"))
def main(argv):
args = utils.parse_args("Train a transformer model")
utils.redirect_log_to_file(args.model_dir)
hparams = create_hparams(args.model_dir, args.configs, initialize=True)
utils.check_git_hash(args.model_dir)
# Prepare data
data.load_vocab(hparams)
train_input_fn = data.InputPipeline(None, None, hparams.record_train_file,
tf.estimator.ModeKeys.TRAIN, hparams)
eval_input_fn = data.InputPipeline(None, None, hparams.record_eval_file,
tf.estimator.ModeKeys.EVAL, hparams)
# Training
log_samples_hook = tf.train.LoggingTensorHook(
['targets', 'predictions'],
at_end=True,
formatter=tensors_to_string(hparams))
train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn,
max_steps=hparams.train_steps)
eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn,
steps=hparams.eval_steps,
hooks=[log_samples_hook])
distribution = tf.contrib.distribute.MirroredStrategy()
run_config = tf.estimator.RunConfig(
model_dir=args.model_dir,
train_distribute=distribution,
save_summary_steps=hparams.save_summary_steps,
save_checkpoints_steps=hparams.save_checkpoints_steps,
keep_checkpoint_max=hparams.n_checkpoints)
estimator = tf.estimator.Estimator(model_fn=model.build_model_fn(hparams),
config=run_config,
model_dir=args.model_dir)
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
def main():
hps = utils.get_hparams()
logger = utils.get_logger(hps.model_dir)
logger.info(hps)
utils.check_git_hash(hps.model_dir)
torch.manual_seed(hps.train.seed)
train_dataset = TextMelLoader(hps.data.training_files, hps.data)
collate_fn = TextMelCollate(1)
train_loader = DataLoader(train_dataset, num_workers=8, shuffle=True,
batch_size=hps.train.batch_size, pin_memory=True,
drop_last=True, collate_fn=collate_fn)
generator = FlowGenerator_DDI(
len(symbols),
out_channels=hps.data.n_mel_channels,
**hps.model).cuda()
optimizer_g = commons.Adam(generator.parameters(), scheduler=hps.train.scheduler, dim_model=hps.model.hidden_channels, warmup_steps=hps.train.warmup_steps, lr=hps.train.learning_rate, betas=hps.train.betas, eps=hps.train.eps)
generator.train()
for batch_idx, (x, x_lengths, y, y_lengths) in enumerate(train_loader):
x, x_lengths = x.cuda(), x_lengths.cuda()
y, y_lengths = y.cuda(), y_lengths.cuda()
_ = generator(x, x_lengths, y, y_lengths, gen=False)
break
# check for pretrained and load it without a an optimizer
pretrained_checkpoint_path = os.path.join(hps.model_dir, "pretrained.pth")
if os.path.isfile(pretrained_checkpoint_path):
logger.info("Loading pretrained checkpoint: %s" % pretrained_checkpoint_path)
model, optimizer, learning_rate, iteration = utils.load_checkpoint(pretrained_checkpoint_path, generator)
utils.save_checkpoint(model, optimizer_g, hps.train.learning_rate, 0, os.path.join(hps.model_dir, "ddi_G.pth"))
else:
utils.save_checkpoint(generator, optimizer_g, hps.train.learning_rate, 0, os.path.join(hps.model_dir, "ddi_G.pth"))
import json
import argparse
import torch
from torch import nn, optim
from torch.nn import functional as F
from torch.utils.data import DataLoader
from data import LJspeechDataset, collate_fn, collate_fn_synthesize
import models
import commons
import utils
hps = utils.get_hparams()
logger = utils.get_logger(hps.model_dir)
logger.info(hps)
utils.check_git_hash(hps.model_dir)
use_cuda = hps.train.use_cuda and torch.cuda.is_available()
torch.manual_seed(hps.train.seed)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 8, 'pin_memory': True} if use_cuda else {}
train_dataset = LJspeechDataset(hps.data.data_path, True, 0.1)
test_dataset = LJspeechDataset(hps.data.data_path, False, 0.1)
train_loader = DataLoader(train_dataset,
batch_size=hps.train.batch_size,
shuffle=True,
collate_fn=collate_fn,
**kwargs)
test_loader = DataLoader(test_dataset,
batch_size=hps.train.batch_size,
def run(rank, n_gpus, hps):
global global_step
if rank == 0:
logger = utils.get_logger(hps.model_dir)
logger.info(hps)
utils.check_git_hash(hps.model_dir)
writer = SummaryWriter(log_dir=hps.model_dir)
writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)
torch.manual_seed(hps.train.seed)
torch.cuda.set_device(rank)
train_dataset = AudioSpecLoader(hps.data.training_files, hps.data)
train_sampler = DistributedBucketSampler(
train_dataset,
hps.train.batch_size,
[32,300,400,500,600,700,800,900,1000],
num_replicas=n_gpus,
rank=rank,
shuffle=True)
collate_fn = AudioSpecCollate()
train_loader = DataLoader(train_dataset, num_workers=8, shuffle=False, pin_memory=True,
collate_fn=collate_fn, batch_sampler=train_sampler)
if rank == 0:
eval_dataset = AudioSpecLoader(hps.data.validation_files, hps.data)
eval_loader = DataLoader(eval_dataset, num_workers=8, shuffle=False,
batch_size=hps.train.batch_size, pin_memory=True,
drop_last=False, collate_fn=collate_fn)
net_g = SynthesizerTrn(
hps.data.filter_length // 2 + 1,
hps.train.segment_size // hps.data.hop_length,
n_speakers=hps.data.n_speakers,
**hps.model).cuda(rank)
net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(rank)
optim_g = torch.optim.AdamW(
net_g.parameters(),
hps.train.learning_rate,
betas=hps.train.betas,
eps=hps.train.eps)
optim_d = torch.optim.AdamW(
net_d.parameters(),
hps.train.learning_rate,
betas=hps.train.betas,
eps=hps.train.eps)
net_g = DDP(net_g, find_unused_parameters=True,device_ids=[rank])
net_d = DDP(net_d, device_ids=[rank])
try:
_, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g)
_, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, optim_d)
global_step = (epoch_str - 1) * len(train_loader)
except:
epoch_str = 1
global_step = 0
scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str-2)
scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str-2)
scaler = GradScaler(enabled=hps.train.fp16_run)
for epoch in range(epoch_str, hps.train.epochs + 1):
if rank==0:
train_and_evaluate(rank, epoch, hps, [net_g, net_d], [optim_g, optim_d], [scheduler_g, scheduler_d], scaler, [train_loader, eval_loader], logger, [writer, writer_eval])
else:
train_and_evaluate(rank, epoch, hps, [net_g, net_d], [optim_g, optim_d], [scheduler_g, scheduler_d], scaler, [train_loader, None], None, None)
scheduler_g.step()
scheduler_d.step()