Medix-AI/coqui-tts
Medix-AI
/
coqui-tts
mirror of https://github.com/coqui-ai/TTS.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix distributed training for train_* scripts

This commit is contained in:
erogol 2020-10-16 17:53:05 +02:00
parent 193b81b273
commit a1582a0e12
7 changed files with 99 additions and 82 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

113
TTS/bin/train_gan_vocoder.py
View File

@ -19,13 +19,16 @@ from TTS.utils.tensorboard_logger import TensorboardLogger
from TTS.utils.training import setup_torch_training_env from TTS.utils.training import setup_torch_training_env
from TTS.vocoder.datasets.gan_dataset import GANDataset from TTS.vocoder.datasets.gan_dataset import GANDataset
from TTS.vocoder.datasets.preprocess import load_wav_data, load_wav_feat_data from TTS.vocoder.datasets.preprocess import load_wav_data, load_wav_feat_data
# from distribute import (DistributedSampler, apply_gradient_allreduce,
# init_distributed, reduce_tensor)
from TTS.vocoder.layers.losses import DiscriminatorLoss, GeneratorLoss from TTS.vocoder.layers.losses import DiscriminatorLoss, GeneratorLoss
from TTS.vocoder.utils.generic_utils import (plot_results, setup_discriminator, from TTS.vocoder.utils.generic_utils import (plot_results, setup_discriminator,
setup_generator) setup_generator)
from TTS.vocoder.utils.io import save_best_model, save_checkpoint from TTS.vocoder.utils.io import save_best_model, save_checkpoint
# DISTRIBUTED
from torch.nn.parallel import DistributedDataParallel as DDP_th
from torch.utils.data.distributed import DistributedSampler
from TTS.utils.distribute import init_distributed
use_cuda, num_gpus = setup_torch_training_env(True, True) use_cuda, num_gpus = setup_torch_training_env(True, True)
@ -45,12 +48,12 @@ def setup_loader(ap, is_val=False, verbose=False):
use_cache=c.use_cache, use_cache=c.use_cache,
verbose=verbose) verbose=verbose)
dataset.shuffle_mapping() dataset.shuffle_mapping()
# sampler = DistributedSampler(dataset) if num_gpus > 1 else None sampler = DistributedSampler(dataset, shuffle=True) if num_gpus > 1 else None
loader = DataLoader(dataset, loader = DataLoader(dataset,
batch_size=1 if is_val else c.batch_size, batch_size=1 if is_val else c.batch_size,
shuffle=True, shuffle=False if num_gpus > 1 else True,
drop_last=False, drop_last=False,
sampler=None, sampler=sampler,
num_workers=c.num_val_loader_workers num_workers=c.num_val_loader_workers
if is_val else c.num_loader_workers, if is_val else c.num_loader_workers,
pin_memory=False) pin_memory=False)
@ -243,41 +246,42 @@ def train(model_G, criterion_G, optimizer_G, model_D, criterion_D, optimizer_D,
c_logger.print_train_step(batch_n_iter, num_iter, global_step, c_logger.print_train_step(batch_n_iter, num_iter, global_step,
log_dict, loss_dict, keep_avg.avg_values) log_dict, loss_dict, keep_avg.avg_values)
# plot step stats if args.rank == 0:
if global_step % 10 == 0: # plot step stats
iter_stats = { if global_step % 10 == 0:
"lr_G": current_lr_G, iter_stats = {
"lr_D": current_lr_D, "lr_G": current_lr_G,
"step_time": step_time "lr_D": current_lr_D,
} "step_time": step_time
iter_stats.update(loss_dict) }
tb_logger.tb_train_iter_stats(global_step, iter_stats) iter_stats.update(loss_dict)
tb_logger.tb_train_iter_stats(global_step, iter_stats)
# save checkpoint # save checkpoint
if global_step % c.save_step == 0: if global_step % c.save_step == 0:
if c.checkpoint: if c.checkpoint:
# save model # save model
save_checkpoint(model_G, save_checkpoint(model_G,
optimizer_G, optimizer_G,
scheduler_G, scheduler_G,
model_D, model_D,
optimizer_D, optimizer_D,
scheduler_D, scheduler_D,
global_step, global_step,
epoch, epoch,
OUT_PATH, OUT_PATH,
model_losses=loss_dict) model_losses=loss_dict)
# compute spectrograms # compute spectrograms
figures = plot_results(y_hat_vis, y_G, ap, global_step, figures = plot_results(y_hat_vis, y_G, ap, global_step,
'train') 'train')
tb_logger.tb_train_figures(global_step, figures) tb_logger.tb_train_figures(global_step, figures)
# Sample audio # Sample audio
sample_voice = y_hat_vis[0].squeeze(0).detach().cpu().numpy() sample_voice = y_hat_vis[0].squeeze(0).detach().cpu().numpy()
tb_logger.tb_train_audios(global_step, tb_logger.tb_train_audios(global_step,
{'train/audio': sample_voice}, {'train/audio': sample_voice},
c.audio["sample_rate"]) c.audio["sample_rate"])
end_time = time.time() end_time = time.time()
# print epoch stats # print epoch stats
@ -286,7 +290,8 @@ def train(model_G, criterion_G, optimizer_G, model_D, criterion_D, optimizer_D,
# Plot Training Epoch Stats # Plot Training Epoch Stats
epoch_stats = {"epoch_time": epoch_time} epoch_stats = {"epoch_time": epoch_time}
epoch_stats.update(keep_avg.avg_values) epoch_stats.update(keep_avg.avg_values)
tb_logger.tb_train_epoch_stats(global_step, epoch_stats) if args.rank == 0:
tb_logger.tb_train_epoch_stats(global_step, epoch_stats)
# TODO: plot model stats # TODO: plot model stats
# if c.tb_model_param_stats: # if c.tb_model_param_stats:
# tb_logger.tb_model_weights(model, global_step) # tb_logger.tb_model_weights(model, global_step)
@ -417,20 +422,21 @@ def evaluate(model_G, criterion_G, model_D, criterion_D, ap, global_step, epoch)
if c.print_eval: if c.print_eval:
c_logger.print_eval_step(num_iter, loss_dict, keep_avg.avg_values) c_logger.print_eval_step(num_iter, loss_dict, keep_avg.avg_values)
# compute spectrograms if args.rank == 0:
figures = plot_results(y_hat, y_G, ap, global_step, 'eval') # compute spectrograms
tb_logger.tb_eval_figures(global_step, figures) figures = plot_results(y_hat, y_G, ap, global_step, 'eval')
tb_logger.tb_eval_figures(global_step, figures)
# Sample audio # Sample audio
sample_voice = y_hat[0].squeeze(0).detach().cpu().numpy() sample_voice = y_hat[0].squeeze(0).detach().cpu().numpy()
tb_logger.tb_eval_audios(global_step, {'eval/audio': sample_voice}, tb_logger.tb_eval_audios(global_step, {'eval/audio': sample_voice},
c.audio["sample_rate"]) c.audio["sample_rate"])
# synthesize a full voice tb_logger.tb_eval_stats(global_step, keep_avg.avg_values)
# synthesize a full voice
data_loader.return_segments = False data_loader.return_segments = False
tb_logger.tb_eval_stats(global_step, keep_avg.avg_values)
return keep_avg.avg_values return keep_avg.avg_values
@ -450,9 +456,9 @@ def main(args): # pylint: disable=redefined-outer-name
ap = AudioProcessor(**c.audio) ap = AudioProcessor(**c.audio)
# DISTRUBUTED # DISTRUBUTED
# if num_gpus > 1: if num_gpus > 1:
# init_distributed(args.rank, num_gpus, args.group_id, init_distributed(args.rank, num_gpus, args.group_id,
# c.distributed["backend"], c.distributed["url"]) c.distributed["backend"], c.distributed["url"])
# setup models # setup models
model_gen = setup_generator(c) model_gen = setup_generator(c)
@ -532,8 +538,9 @@ def main(args): # pylint: disable=redefined-outer-name
criterion_disc.cuda() criterion_disc.cuda()
# DISTRUBUTED # DISTRUBUTED
# if num_gpus > 1: if num_gpus > 1:
# model = apply_gradient_allreduce(model) model_gen = DDP_th(model_gen, device_ids=[args.rank])
model_disc = DDP_th(model_disc, device_ids=[args.rank])
num_params = count_parameters(model_gen) num_params = count_parameters(model_gen)
print(" > Generator has {} parameters".format(num_params), flush=True) print(" > Generator has {} parameters".format(num_params), flush=True)

16
TTS/bin/train_glow_tts.py
View File

@ -11,6 +11,7 @@ import traceback
import torch import torch
from random import randrange from random import randrange
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
from TTS.tts.datasets.preprocess import load_meta_data from TTS.tts.datasets.preprocess import load_meta_data
from TTS.tts.datasets.TTSDataset import MyDataset from TTS.tts.datasets.TTSDataset import MyDataset
from TTS.tts.layers.losses import GlowTTSLoss from TTS.tts.layers.losses import GlowTTSLoss
@ -34,6 +35,13 @@ from TTS.utils.tensorboard_logger import TensorboardLogger
from TTS.utils.training import (NoamLR, check_update, from TTS.utils.training import (NoamLR, check_update,
setup_torch_training_env) setup_torch_training_env)
# DISTRIBUTED
from apex.parallel import DistributedDataParallel as DDP_apex
from torch.nn.parallel import DistributedDataParallel as DDP_th
from torch.utils.data.distributed import DistributedSampler
from TTS.utils.distribute import init_distributed, reduce_tensor
use_cuda, num_gpus = setup_torch_training_env(True, False) use_cuda, num_gpus = setup_torch_training_env(True, False)
def setup_loader(ap, r, is_val=False, verbose=False, speaker_mapping=None): def setup_loader(ap, r, is_val=False, verbose=False, speaker_mapping=None):
@ -481,10 +489,9 @@ def main(args): # pylint: disable=redefined-outer-name
optimizer = RAdam(model.parameters(), lr=c.lr, weight_decay=0, betas=(0.9, 0.98), eps=1e-9) optimizer = RAdam(model.parameters(), lr=c.lr, weight_decay=0, betas=(0.9, 0.98), eps=1e-9)
criterion = GlowTTSLoss() criterion = GlowTTSLoss()
if c.apex_amp_level: if c.apex_amp_level is not None:
# pylint: disable=import-outside-toplevel # pylint: disable=import-outside-toplevel
from apex import amp from apex import amp
from apex.parallel import DistributedDataParallel as DDP
model.cuda() model.cuda()
model, optimizer = amp.initialize(model, optimizer, opt_level=c.apex_amp_level) model, optimizer = amp.initialize(model, optimizer, opt_level=c.apex_amp_level)
else: else:
@ -523,7 +530,10 @@ def main(args): # pylint: disable=redefined-outer-name
# DISTRUBUTED # DISTRUBUTED
if num_gpus > 1: if num_gpus > 1:
model = DDP(model) if c.apex_amp_level is not None:
model = DDP_apex(model)
else:
model = DDP_th(model, device_ids=[args.rank])
if c.noam_schedule: if c.noam_schedule:
scheduler = NoamLR(optimizer, scheduler = NoamLR(optimizer,

2
TTS/bin/train_tts.py
View File

@ -38,8 +38,10 @@ from TTS.utils.training import (NoamLR, adam_weight_decay, check_update,
gradual_training_scheduler, set_weight_decay, gradual_training_scheduler, set_weight_decay,
setup_torch_training_env) setup_torch_training_env)
use_cuda, num_gpus = setup_torch_training_env(True, False) use_cuda, num_gpus = setup_torch_training_env(True, False)
def setup_loader(ap, r, is_val=False, verbose=False, speaker_mapping=None): def setup_loader(ap, r, is_val=False, verbose=False, speaker_mapping=None):
if is_val and not c.run_eval: if is_val and not c.run_eval:
loader = None loader = None

33
TTS/bin/train_wavegrad.py
View File

@ -4,12 +4,9 @@ import os
import sys import sys
import time import time
import traceback import traceback
from inspect import signature
import torch import torch
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from TTS.utils.audio import AudioProcessor from TTS.utils.audio import AudioProcessor
from TTS.utils.console_logger import ConsoleLogger from TTS.utils.console_logger import ConsoleLogger
@ -20,14 +17,18 @@ from TTS.utils.io import copy_config_file, load_config
from TTS.utils.radam import RAdam from TTS.utils.radam import RAdam
from TTS.utils.tensorboard_logger import TensorboardLogger from TTS.utils.tensorboard_logger import TensorboardLogger
from TTS.utils.training import setup_torch_training_env from TTS.utils.training import setup_torch_training_env
from TTS.vocoder.datasets.wavegrad_dataset import WaveGradDataset
from TTS.vocoder.datasets.preprocess import load_wav_data, load_wav_feat_data from TTS.vocoder.datasets.preprocess import load_wav_data, load_wav_feat_data
from TTS.utils.distribute import init_distributed, reduce_tensor from TTS.vocoder.datasets.wavegrad_dataset import WaveGradDataset
from TTS.vocoder.layers.losses import DiscriminatorLoss, GeneratorLoss from TTS.vocoder.utils.generic_utils import plot_results, setup_generator
from TTS.vocoder.utils.generic_utils import (plot_results, setup_discriminator,
setup_generator)
from TTS.vocoder.utils.io import save_best_model, save_checkpoint from TTS.vocoder.utils.io import save_best_model, save_checkpoint
# DISTRIBUTED
from apex.parallel import DistributedDataParallel as DDP_apex
from torch.nn.parallel import DistributedDataParallel as DDP_th
from torch.utils.data.distributed import DistributedSampler
from TTS.utils.distribute import init_distributed
use_cuda, num_gpus = setup_torch_training_env(True, True) use_cuda, num_gpus = setup_torch_training_env(True, True)
@ -111,11 +112,6 @@ def train(model, criterion, optimizer,
else: else:
loss.backward() loss.backward()
if amp:
amp_opt_params = amp.master_params(optimizer)
else:
amp_opt_params = None
if c.clip_grad > 0: if c.clip_grad > 0:
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(),
c.clip_grad) c.clip_grad)
@ -279,7 +275,6 @@ def evaluate(model, criterion, ap, global_step, epoch):
return keep_avg.avg_values return keep_avg.avg_values
# FIXME: move args definition/parsing inside of main?
def main(args): # pylint: disable=redefined-outer-name def main(args): # pylint: disable=redefined-outer-name
# pylint: disable=global-variable-undefined # pylint: disable=global-variable-undefined
global train_data, eval_data global train_data, eval_data
@ -305,10 +300,9 @@ def main(args): # pylint: disable=redefined-outer-name
optimizer = RAdam(model.parameters(), lr=c.lr, weight_decay=0) optimizer = RAdam(model.parameters(), lr=c.lr, weight_decay=0)
# DISTRIBUTED # DISTRIBUTED
if c.apex_amp_level: if c.apex_amp_level is not None:
# pylint: disable=import-outside-toplevel # pylint: disable=import-outside-toplevel
from apex import amp from apex import amp
from apex.parallel import DistributedDataParallel as DDP
model.cuda() model.cuda()
model, optimizer = amp.initialize(model, optimizer, opt_level=c.apex_amp_level) model, optimizer = amp.initialize(model, optimizer, opt_level=c.apex_amp_level)
else: else:
@ -363,7 +357,10 @@ def main(args): # pylint: disable=redefined-outer-name
# DISTRUBUTED # DISTRUBUTED
if num_gpus > 1: if num_gpus > 1:
model = DDP(model) if c.apex_amp_level is not None:
model = DDP_apex(model)
else:
model = DDP_th(model, device_ids=[args.rank])
num_params = count_parameters(model) num_params = count_parameters(model)
print(" > WaveGrad has {} parameters".format(num_params), flush=True) print(" > WaveGrad has {} parameters".format(num_params), flush=True)
@ -447,7 +444,7 @@ if __name__ == '__main__':
_ = os.path.dirname(os.path.realpath(__file__)) _ = os.path.dirname(os.path.realpath(__file__))
# DISTRIBUTED # DISTRIBUTED
if c.apex_amp_level: if c.apex_amp_level is not None:
print(" > apex AMP level: ", c.apex_amp_level) print(" > apex AMP level: ", c.apex_amp_level)
OUT_PATH = args.continue_path OUT_PATH = args.continue_path

3
TTS/tts/configs/glow_tts_gated_conv.json
View File

@ -54,9 +54,10 @@
"add_blank": false, // if true add a new token after each token of the sentence. This increases the size of the input sequence, but has considerably improved the prosody of the GlowTTS model. "add_blank": false, // if true add a new token after each token of the sentence. This increases the size of the input sequence, but has considerably improved the prosody of the GlowTTS model.
// DISTRIBUTED TRAINING // DISTRIBUTED TRAINING
"apex_amp_level": null, // APEX amp optimization level. "O1" is currently supported.
"distributed":{ "distributed":{
"backend": "nccl", "backend": "nccl",
"url": "tcp:\/\/localhost:54321" "url": "tcp:\/\/localhost:54323"
}, },
"reinit_layers": [], // give a list of layer names to restore from the given checkpoint. If not defined, it reloads all heuristically matching layers. "reinit_layers": [], // give a list of layer names to restore from the given checkpoint. If not defined, it reloads all heuristically matching layers.

8
TTS/vocoder/configs/universal_fullband_melgan.json
View File

@ -31,13 +31,13 @@
"symmetric_norm": true, // move normalization to range [-1, 1] "symmetric_norm": true, // move normalization to range [-1, 1]
"max_norm": 4.0, // scale normalization to range [-max_norm, max_norm] or [0, max_norm] "max_norm": 4.0, // scale normalization to range [-max_norm, max_norm] or [0, max_norm]
"clip_norm": true, // clip normalized values into the range. "clip_norm": true, // clip normalized values into the range.
"stats_path": "/data/rw/home/Data/LibriTTS/scale_stats.npy" // DO NOT USE WITH MULTI_SPEAKER MODEL. scaler stats file computed by 'compute_statistics.py'. If it is defined, mean-std based notmalization is used and other normalization params are ignored "stats_path": "/home/erogol/Data/libritts/LibriTTS/scale_stats.npy" // DO NOT USE WITH MULTI_SPEAKER MODEL. scaler stats file computed by 'compute_statistics.py'. If it is defined, mean-std based notmalization is used and other normalization params are ignored
}, },
// DISTRIBUTED TRAINING // DISTRIBUTED TRAINING
"distributed":{ "distributed":{
"backend": "nccl", "backend": "nccl",
"url": "tcp:\/\/localhost:54321" "url": "tcp:\/\/localhost:54324"
}, },
// MODEL PARAMETERS // MODEL PARAMETERS
@ -83,7 +83,7 @@
}, },
// DATASET // DATASET
"data_path": "/data5/rw/home/Data/LibriTTS/LibriTTS/train-clean-360/", "data_path": "/home/erogol/Data/libritts/LibriTTS/train-clean-360/",
"feature_path": null, "feature_path": null,
"seq_len": 16384, "seq_len": 16384,
"pad_short": 2000, "pad_short": 2000,
@ -132,7 +132,7 @@
"eval_split_size": 10, "eval_split_size": 10,
// PATHS // PATHS
"output_path": "/data4/rw/home/Trainings/LJSpeech/" "output_path": "/home/erogol/Models/"
} }

6
TTS/vocoder/configs/wavegrad_libritts.json
View File

@ -34,10 +34,10 @@
}, },
// DISTRIBUTED TRAINING // DISTRIBUTED TRAINING
"apex_amp_level": "O1", // amp optimization level. "O1" is currentl supported. "apex_amp_level": null, // APEX amp optimization level. "O1" is currently supported.
"distributed":{ "distributed":{
"backend": "nccl", "backend": "nccl",
"url": "tcp:\/\/localhost:54321" "url": "tcp:\/\/localhost:54322"
}, },
"target_loss": "avg_wavegrad_loss", // loss value to pick the best model to save after each epoch "target_loss": "avg_wavegrad_loss", // loss value to pick the best model to save after each epoch
@ -47,7 +47,7 @@
"model_params":{ "model_params":{
"x_conv_channels":32, "x_conv_channels":32,
"c_conv_channels":768, "c_conv_channels":768,
"ublock_out_channels": [768, 512, 512, 256, 128], "ublock_out_channels": [512, 512, 256, 128, 128],
"dblock_out_channels": [128, 128, 256, 512], "dblock_out_channels": [128, 128, 256, 512],
"upsample_factors": [4, 4, 4, 2, 2], "upsample_factors": [4, 4, 4, 2, 2],
"upsample_dilations": [ "upsample_dilations": [