def __init__(self, config, train: DataLoader, val: DataLoader):
self.config = config
self.train_dataset = train
self.val_dataset = val
self.adv_lambda = config['model']['adv_lambda']
self.metric_counter = MetricCounter(config['experiment_desc'])
self.warmup_epochs = config['warmup_num']
def __init__(self, config, train: DataLoader, val: DataLoader):
self.config = config
self.train_dataset = train
self.val_dataset = val
self.adv_lambda = config['model']['adv_lambda']
self.metric_counter = MetricCounter(config['experiment_desc'])
self.warmup_epochs = config['warmup_num']
gpu_id = self.config['gpu_id']
self.device = torch.device('cuda:{}'.format(gpu_id) if (torch.cuda.is_available() and gpu_id > 0) else "cpu")
def __init__(self, config):
self.config = config
self.train_dataset = self._get_dataset(config, 'train')
self.val_dataset = self._get_dataset(config, 'test')
self.experiment_name = config['experiment_desc'] + '_' + config[
'model']['g_name']
self.experiment_name += '_content' + str(config['model']['content_coef'])+ '_feature' \
+ str(config['model']['feature_coef']) + '_adv' \
+ str(config['model']['adv_coef'])
self.metric_counter = MetricCounter(self.experiment_name)
self.warmup_epochs = config['warmup_num']
def __init__(self, config):
self.config = config
self.dataset = self._get_dataset(config["dataroot"],
config["seq_length"])
dlen = len(self.dataset)
splitlen = [int(0.8 * dlen), int(0.1 * dlen), 0]
splitlen[2] = dlen - sum(splitlen)
self.train_dataset, self.val_dataset, self.test_dataset = data.random_split(
self.dataset, splitlen)
self.experiment_name = f"{config['experiment_desc']}_{config['model']['model_n']}"
self.metric_counter = MetricCounter(self.experiment_name,
self.config["print_every"])
def train():
prepare_directories()
metric_counter = MetricCounter(exp_name=LOG_DIR)
parameters = dict(audio_configs)
parameters["text_cleaner"] = configs["text_cleaner"]
parameters["outputs_per_step"] = configs["r"]
train_dataset = TextSpeechDataset(root_dir=data_configs["data_path"],
annotations_file=data_configs["annotations_train"], parameters=parameters)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True,
collate_fn=train_dataset.collate_fn,
num_workers=num_workers, drop_last=False, pin_memory=True)
val_dataset = TextSpeechDataset(root_dir=data_configs["data_path"],
annotations_file=data_configs["annotations_val"], parameters=parameters)
val_loader = DataLoader(dataset=val_dataset, batch_size=eval_batch_size, num_workers=num_workers,
collate_fn=val_dataset.collate_fn, drop_last=False, pin_memory=True)
model = Tacotron(embedding_dim=configs.pop("embedding_size"),
linear_dim=audio_configs["frequency"],
mel_dim=audio_configs["mels_size"],
r=configs.pop("r"))
if use_cuda:
model = torch.nn.DataParallel(model.to("cuda"))
optimizer = optim.Adam(params=model.parameters(), lr=train_configs["lr"])
criterion = L1LossMasked()
if args.resume:
model.load_state_dict(torch.load(args.resume))
n_priority_freq = int(3000 / (audio_configs["sample_rate"] * 0.5) * audio_configs["frequency"])
for epoch in range(train_configs["epochs"]):
audio_signal = run_epoch(model, train_loader, optimizer, criterion, metric_counter, epoch, n_priority_freq)
run_validate(model, val_loader, criterion, metric_counter, n_priority_freq)
if metric_counter.update_best_model():
torch.save(model.state_dict(), os.path.join(os.path.join(WEIGHTS_SAVE_PATH,
f"best_{configs['experiment_name']}.pth.tar")))
audio_signal = train_dataset.ap.spectrogram_to_wav(audio_signal.T)
metric_counter.write_audio_to_tensorboard("Audio", audio_signal, epoch, audio_configs["sample_rate"])
torch.save(model.state_dict(), os.path.join(os.path.join(WEIGHTS_SAVE_PATH,
f"last_{configs['experiment_name']}.pth.tar")))
print(metric_counter.loss_message())
logging.debug(
f"Experiment Name: {configs['experiment_name']}, Epoch: {epoch}, Loss: {metric_counter.loss_message()}")
def __init__(self, config, val: DataLoader):
self.config = config
self.val_dataset = val
self.metric_counter = MetricCounter()