styletts2 / train_first.py
mrfakename's picture
Initial Commit
635f007
raw
history blame
19 kB
import os
import os.path as osp
import re
import sys
import yaml
import shutil
import numpy as np
import torch
import click
import warnings
warnings.simplefilter("ignore")
# load packages
import random
import yaml
from munch import Munch
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torchaudio
import librosa
from models import *
from meldataset import build_dataloader
from utils import *
from losses import *
from optimizers import build_optimizer
import time
from accelerate import Accelerator
from accelerate.utils import LoggerType
from accelerate import DistributedDataParallelKwargs
from torch.utils.tensorboard import SummaryWriter
import logging
from accelerate.logging import get_logger
logger = get_logger(__name__, log_level="DEBUG")
@click.command()
@click.option("-p", "--config_path", default="Configs/config.yml", type=str)
def main(config_path):
config = yaml.safe_load(open(config_path))
log_dir = config["log_dir"]
if not osp.exists(log_dir):
os.makedirs(log_dir, exist_ok=True)
shutil.copy(config_path, osp.join(log_dir, osp.basename(config_path)))
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
accelerator = Accelerator(
project_dir=log_dir, split_batches=True, kwargs_handlers=[ddp_kwargs]
)
if accelerator.is_main_process:
writer = SummaryWriter(log_dir + "/tensorboard")
# write logs
file_handler = logging.FileHandler(osp.join(log_dir, "train.log"))
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(
logging.Formatter("%(levelname)s:%(asctime)s: %(message)s")
)
logger.logger.addHandler(file_handler)
batch_size = config.get("batch_size", 10)
device = accelerator.device
epochs = config.get("epochs_1st", 200)
save_freq = config.get("save_freq", 2)
log_interval = config.get("log_interval", 10)
saving_epoch = config.get("save_freq", 2)
data_params = config.get("data_params", None)
sr = config["preprocess_params"].get("sr", 24000)
train_path = data_params["train_data"]
val_path = data_params["val_data"]
root_path = data_params["root_path"]
min_length = data_params["min_length"]
OOD_data = data_params["OOD_data"]
max_len = config.get("max_len", 200)
# load data
train_list, val_list = get_data_path_list(train_path, val_path)
train_dataloader = build_dataloader(
train_list,
root_path,
OOD_data=OOD_data,
min_length=min_length,
batch_size=batch_size,
num_workers=2,
dataset_config={},
device=device,
)
val_dataloader = build_dataloader(
val_list,
root_path,
OOD_data=OOD_data,
min_length=min_length,
batch_size=batch_size,
validation=True,
num_workers=0,
device=device,
dataset_config={},
)
with accelerator.main_process_first():
# load pretrained ASR model
ASR_config = config.get("ASR_config", False)
ASR_path = config.get("ASR_path", False)
text_aligner = load_ASR_models(ASR_path, ASR_config)
# load pretrained F0 model
F0_path = config.get("F0_path", False)
pitch_extractor = load_F0_models(F0_path)
# load BERT model
from Utils.PLBERT.util import load_plbert
BERT_path = config.get("PLBERT_dir", False)
plbert = load_plbert(BERT_path)
scheduler_params = {
"max_lr": float(config["optimizer_params"].get("lr", 1e-4)),
"pct_start": float(config["optimizer_params"].get("pct_start", 0.0)),
"epochs": epochs,
"steps_per_epoch": len(train_dataloader),
}
model_params = recursive_munch(config["model_params"])
multispeaker = model_params.multispeaker
model = build_model(model_params, text_aligner, pitch_extractor, plbert)
best_loss = float("inf") # best test loss
loss_train_record = list([])
loss_test_record = list([])
loss_params = Munch(config["loss_params"])
TMA_epoch = loss_params.TMA_epoch
for k in model:
model[k] = accelerator.prepare(model[k])
train_dataloader, val_dataloader = accelerator.prepare(
train_dataloader, val_dataloader
)
_ = [model[key].to(device) for key in model]
# initialize optimizers after preparing models for compatibility with FSDP
optimizer = build_optimizer(
{key: model[key].parameters() for key in model},
scheduler_params_dict={key: scheduler_params.copy() for key in model},
lr=float(config["optimizer_params"].get("lr", 1e-4)),
)
for k, v in optimizer.optimizers.items():
optimizer.optimizers[k] = accelerator.prepare(optimizer.optimizers[k])
optimizer.schedulers[k] = accelerator.prepare(optimizer.schedulers[k])
with accelerator.main_process_first():
if config.get("pretrained_model", "") != "":
model, optimizer, start_epoch, iters = load_checkpoint(
model,
optimizer,
config["pretrained_model"],
load_only_params=config.get("load_only_params", True),
)
else:
start_epoch = 0
iters = 0
# in case not distributed
try:
n_down = model.text_aligner.module.n_down
except:
n_down = model.text_aligner.n_down
# wrapped losses for compatibility with mixed precision
stft_loss = MultiResolutionSTFTLoss().to(device)
gl = GeneratorLoss(model.mpd, model.msd).to(device)
dl = DiscriminatorLoss(model.mpd, model.msd).to(device)
wl = WavLMLoss(model_params.slm.model, model.wd, sr, model_params.slm.sr).to(device)
for epoch in range(start_epoch, epochs):
running_loss = 0
start_time = time.time()
_ = [model[key].train() for key in model]
for i, batch in enumerate(train_dataloader):
waves = batch[0]
batch = [b.to(device) for b in batch[1:]]
texts, input_lengths, _, _, mels, mel_input_length, _ = batch
with torch.no_grad():
mask = length_to_mask(mel_input_length // (2**n_down)).to("cuda")
text_mask = length_to_mask(input_lengths).to(texts.device)
ppgs, s2s_pred, s2s_attn = model.text_aligner(mels, mask, texts)
s2s_attn = s2s_attn.transpose(-1, -2)
s2s_attn = s2s_attn[..., 1:]
s2s_attn = s2s_attn.transpose(-1, -2)
with torch.no_grad():
attn_mask = (
(~mask)
.unsqueeze(-1)
.expand(mask.shape[0], mask.shape[1], text_mask.shape[-1])
.float()
.transpose(-1, -2)
)
attn_mask = (
attn_mask.float()
* (~text_mask)
.unsqueeze(-1)
.expand(text_mask.shape[0], text_mask.shape[1], mask.shape[-1])
.float()
)
attn_mask = attn_mask < 1
s2s_attn.masked_fill_(attn_mask, 0.0)
with torch.no_grad():
mask_ST = mask_from_lens(
s2s_attn, input_lengths, mel_input_length // (2**n_down)
)
s2s_attn_mono = maximum_path(s2s_attn, mask_ST)
# encode
t_en = model.text_encoder(texts, input_lengths, text_mask)
# 50% of chance of using monotonic version
if bool(random.getrandbits(1)):
asr = t_en @ s2s_attn
else:
asr = t_en @ s2s_attn_mono
# get clips
mel_input_length_all = accelerator.gather(
mel_input_length
) # for balanced load
mel_len = min(
[int(mel_input_length_all.min().item() / 2 - 1), max_len // 2]
)
mel_len_st = int(mel_input_length.min().item() / 2 - 1)
en = []
gt = []
wav = []
st = []
for bib in range(len(mel_input_length)):
mel_length = int(mel_input_length[bib].item() / 2)
random_start = np.random.randint(0, mel_length - mel_len)
en.append(asr[bib, :, random_start : random_start + mel_len])
gt.append(
mels[bib, :, (random_start * 2) : ((random_start + mel_len) * 2)]
)
y = waves[bib][
(random_start * 2) * 300 : ((random_start + mel_len) * 2) * 300
]
wav.append(torch.from_numpy(y).to(device))
# style reference (better to be different from the GT)
random_start = np.random.randint(0, mel_length - mel_len_st)
st.append(
mels[bib, :, (random_start * 2) : ((random_start + mel_len_st) * 2)]
)
en = torch.stack(en)
gt = torch.stack(gt).detach()
st = torch.stack(st).detach()
wav = torch.stack(wav).float().detach()
# clip too short to be used by the style encoder
if gt.shape[-1] < 80:
continue
with torch.no_grad():
real_norm = log_norm(gt.unsqueeze(1)).squeeze(1).detach()
F0_real, _, _ = model.pitch_extractor(gt.unsqueeze(1))
s = model.style_encoder(
st.unsqueeze(1) if multispeaker else gt.unsqueeze(1)
)
y_rec = model.decoder(en, F0_real, real_norm, s)
# discriminator loss
if epoch >= TMA_epoch:
optimizer.zero_grad()
d_loss = dl(wav.detach().unsqueeze(1).float(), y_rec.detach()).mean()
accelerator.backward(d_loss)
optimizer.step("msd")
optimizer.step("mpd")
else:
d_loss = 0
# generator loss
optimizer.zero_grad()
loss_mel = stft_loss(y_rec.squeeze(), wav.detach())
if epoch >= TMA_epoch: # start TMA training
loss_s2s = 0
for _s2s_pred, _text_input, _text_length in zip(
s2s_pred, texts, input_lengths
):
loss_s2s += F.cross_entropy(
_s2s_pred[:_text_length], _text_input[:_text_length]
)
loss_s2s /= texts.size(0)
loss_mono = F.l1_loss(s2s_attn, s2s_attn_mono) * 10
loss_gen_all = gl(wav.detach().unsqueeze(1).float(), y_rec).mean()
loss_slm = wl(wav.detach(), y_rec).mean()
g_loss = (
loss_params.lambda_mel * loss_mel
+ loss_params.lambda_mono * loss_mono
+ loss_params.lambda_s2s * loss_s2s
+ loss_params.lambda_gen * loss_gen_all
+ loss_params.lambda_slm * loss_slm
)
else:
loss_s2s = 0
loss_mono = 0
loss_gen_all = 0
loss_slm = 0
g_loss = loss_mel
running_loss += accelerator.gather(loss_mel).mean().item()
accelerator.backward(g_loss)
optimizer.step("text_encoder")
optimizer.step("style_encoder")
optimizer.step("decoder")
if epoch >= TMA_epoch:
optimizer.step("text_aligner")
optimizer.step("pitch_extractor")
iters = iters + 1
if (i + 1) % log_interval == 0 and accelerator.is_main_process:
log_print(
"Epoch [%d/%d], Step [%d/%d], Mel Loss: %.5f, Gen Loss: %.5f, Disc Loss: %.5f, Mono Loss: %.5f, S2S Loss: %.5f, SLM Loss: %.5f"
% (
epoch + 1,
epochs,
i + 1,
len(train_list) // batch_size,
running_loss / log_interval,
loss_gen_all,
d_loss,
loss_mono,
loss_s2s,
loss_slm,
),
logger,
)
writer.add_scalar("train/mel_loss", running_loss / log_interval, iters)
writer.add_scalar("train/gen_loss", loss_gen_all, iters)
writer.add_scalar("train/d_loss", d_loss, iters)
writer.add_scalar("train/mono_loss", loss_mono, iters)
writer.add_scalar("train/s2s_loss", loss_s2s, iters)
writer.add_scalar("train/slm_loss", loss_slm, iters)
running_loss = 0
print("Time elasped:", time.time() - start_time)
loss_test = 0
_ = [model[key].eval() for key in model]
with torch.no_grad():
iters_test = 0
for batch_idx, batch in enumerate(val_dataloader):
optimizer.zero_grad()
waves = batch[0]
batch = [b.to(device) for b in batch[1:]]
texts, input_lengths, _, _, mels, mel_input_length, _ = batch
with torch.no_grad():
mask = length_to_mask(mel_input_length // (2**n_down)).to("cuda")
ppgs, s2s_pred, s2s_attn = model.text_aligner(mels, mask, texts)
s2s_attn = s2s_attn.transpose(-1, -2)
s2s_attn = s2s_attn[..., 1:]
s2s_attn = s2s_attn.transpose(-1, -2)
text_mask = length_to_mask(input_lengths).to(texts.device)
attn_mask = (
(~mask)
.unsqueeze(-1)
.expand(mask.shape[0], mask.shape[1], text_mask.shape[-1])
.float()
.transpose(-1, -2)
)
attn_mask = (
attn_mask.float()
* (~text_mask)
.unsqueeze(-1)
.expand(text_mask.shape[0], text_mask.shape[1], mask.shape[-1])
.float()
)
attn_mask = attn_mask < 1
s2s_attn.masked_fill_(attn_mask, 0.0)
# encode
t_en = model.text_encoder(texts, input_lengths, text_mask)
asr = t_en @ s2s_attn
# get clips
mel_input_length_all = accelerator.gather(
mel_input_length
) # for balanced load
mel_len = min(
[int(mel_input_length.min().item() / 2 - 1), max_len // 2]
)
en = []
gt = []
wav = []
for bib in range(len(mel_input_length)):
mel_length = int(mel_input_length[bib].item() / 2)
random_start = np.random.randint(0, mel_length - mel_len)
en.append(asr[bib, :, random_start : random_start + mel_len])
gt.append(
mels[
bib, :, (random_start * 2) : ((random_start + mel_len) * 2)
]
)
y = waves[bib][
(random_start * 2) * 300 : ((random_start + mel_len) * 2) * 300
]
wav.append(torch.from_numpy(y).to("cuda"))
wav = torch.stack(wav).float().detach()
en = torch.stack(en)
gt = torch.stack(gt).detach()
F0_real, _, F0 = model.pitch_extractor(gt.unsqueeze(1))
s = model.style_encoder(gt.unsqueeze(1))
real_norm = log_norm(gt.unsqueeze(1)).squeeze(1)
y_rec = model.decoder(en, F0_real, real_norm, s)
loss_mel = stft_loss(y_rec.squeeze(), wav.detach())
loss_test += accelerator.gather(loss_mel).mean().item()
iters_test += 1
if accelerator.is_main_process:
print("Epochs:", epoch + 1)
log_print(
"Validation loss: %.3f" % (loss_test / iters_test) + "\n\n\n\n", logger
)
print("\n\n\n")
writer.add_scalar("eval/mel_loss", loss_test / iters_test, epoch + 1)
attn_image = get_image(s2s_attn[0].cpu().numpy().squeeze())
writer.add_figure("eval/attn", attn_image, epoch)
with torch.no_grad():
for bib in range(len(asr)):
mel_length = int(mel_input_length[bib].item())
gt = mels[bib, :, :mel_length].unsqueeze(0)
en = asr[bib, :, : mel_length // 2].unsqueeze(0)
F0_real, _, _ = model.pitch_extractor(gt.unsqueeze(1))
F0_real = F0_real.unsqueeze(0)
s = model.style_encoder(gt.unsqueeze(1))
real_norm = log_norm(gt.unsqueeze(1)).squeeze(1)
y_rec = model.decoder(en, F0_real, real_norm, s)
writer.add_audio(
"eval/y" + str(bib),
y_rec.cpu().numpy().squeeze(),
epoch,
sample_rate=sr,
)
if epoch == 0:
writer.add_audio(
"gt/y" + str(bib),
waves[bib].squeeze(),
epoch,
sample_rate=sr,
)
if bib >= 6:
break
if epoch % saving_epoch == 0:
if (loss_test / iters_test) < best_loss:
best_loss = loss_test / iters_test
print("Saving..")
state = {
"net": {key: model[key].state_dict() for key in model},
"optimizer": optimizer.state_dict(),
"iters": iters,
"val_loss": loss_test / iters_test,
"epoch": epoch,
}
save_path = osp.join(log_dir, "epoch_1st_%05d.pth" % epoch)
torch.save(state, save_path)
if accelerator.is_main_process:
print("Saving..")
state = {
"net": {key: model[key].state_dict() for key in model},
"optimizer": optimizer.state_dict(),
"iters": iters,
"val_loss": loss_test / iters_test,
"epoch": epoch,
}
save_path = osp.join(log_dir, config.get("first_stage_path", "first_stage.pth"))
torch.save(state, save_path)
if __name__ == "__main__":
main()