fix
This commit is contained in:
tuna2134
91
train_ms.py
91
train_ms.py
@@ -28,6 +28,7 @@ from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
|
||||
from style_bert_vits2.logging import logger
|
||||
from style_bert_vits2.models import commons, utils
|
||||
from style_bert_vits2.models.hyper_parameters import HyperParameters
|
||||
from style_bert_vits2.models.matcha_flow import configure_matcha_only_training
|
||||
from style_bert_vits2.models.models import (
|
||||
DurationDiscriminator,
|
||||
MultiPeriodDiscriminator,
|
||||
@@ -42,8 +43,7 @@ torch.backends.cudnn.allow_tf32 = (
|
||||
True # If encontered training problem,please try to disable TF32.
|
||||
)
|
||||
torch.set_float32_matmul_precision("medium")
|
||||
torch.backends.cuda.sdp_kernel("flash")
|
||||
torch.backends.cuda.enable_flash_sdp(True)
|
||||
torch.backends.cuda.enable_flash_sdp(False)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(
|
||||
True
|
||||
) # Not available if torch version is lower than 2.0
|
||||
@@ -286,7 +286,7 @@ def run():
|
||||
logger.info("Using normal MAS for VITS1")
|
||||
mas_noise_scale_initial = 0.0
|
||||
noise_scale_delta = 0.0
|
||||
if hps.model.use_duration_discriminator is True:
|
||||
if hps.model.use_duration_discriminator is True and not hps.train.matcha_only:
|
||||
logger.info("Using duration discriminator for VITS2")
|
||||
net_dur_disc = DurationDiscriminator(
|
||||
hps.model.hidden_channels,
|
||||
@@ -322,6 +322,7 @@ def run():
|
||||
matcha_dropout=hps.model.matcha_dropout,
|
||||
matcha_sigma_min=hps.model.matcha_sigma_min,
|
||||
matcha_n_timesteps=hps.model.matcha_n_timesteps,
|
||||
matcha_use_diff_attention=hps.model.matcha_use_diff_attention,
|
||||
inter_channels=hps.model.inter_channels,
|
||||
hidden_channels=hps.model.hidden_channels,
|
||||
filter_channels=hps.model.filter_channels,
|
||||
@@ -365,6 +366,10 @@ def run():
|
||||
for param in net_g.dec.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
if hps.train.matcha_only:
|
||||
trainable = configure_matcha_only_training(net_g)
|
||||
logger.info(f"Matcha-only training enabled ({trainable:,} trainable parameters)")
|
||||
|
||||
net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(local_rank)
|
||||
optim_g = torch.optim.AdamW(
|
||||
filter(lambda p: p.requires_grad, net_g.parameters()),
|
||||
@@ -401,7 +406,7 @@ def run():
|
||||
utils.checkpoints.get_latest_checkpoint_path(model_dir, "DUR_*.pth"),
|
||||
net_dur_disc,
|
||||
optim_dur_disc,
|
||||
skip_optimizer=hps.train.skip_optimizer,
|
||||
skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
|
||||
)
|
||||
if not optim_dur_disc.param_groups[0].get("initial_lr"):
|
||||
optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
|
||||
@@ -409,13 +414,13 @@ def run():
|
||||
utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"),
|
||||
net_g,
|
||||
optim_g,
|
||||
skip_optimizer=hps.train.skip_optimizer,
|
||||
skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
|
||||
)
|
||||
_, optim_d, d_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
|
||||
utils.checkpoints.get_latest_checkpoint_path(model_dir, "D_*.pth"),
|
||||
net_d,
|
||||
optim_d,
|
||||
skip_optimizer=hps.train.skip_optimizer,
|
||||
skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
|
||||
)
|
||||
if not optim_g.param_groups[0].get("initial_lr"):
|
||||
optim_g.param_groups[0]["initial_lr"] = g_resume_lr
|
||||
@@ -618,10 +623,17 @@ def train_and_evaluate(
|
||||
train_loader.batch_sampler.set_epoch(epoch)
|
||||
global global_step
|
||||
|
||||
net_g.train()
|
||||
net_d.train()
|
||||
if net_dur_disc is not None:
|
||||
net_dur_disc.train()
|
||||
if hps.train.matcha_only:
|
||||
net_g.eval()
|
||||
net_g.module.matcha.train()
|
||||
net_d.eval()
|
||||
if net_dur_disc is not None:
|
||||
net_dur_disc.eval()
|
||||
else:
|
||||
net_g.train()
|
||||
net_d.train()
|
||||
if net_dur_disc is not None:
|
||||
net_dur_disc.train()
|
||||
for batch_idx, (
|
||||
x,
|
||||
x_lengths,
|
||||
@@ -637,7 +649,7 @@ def train_and_evaluate(
|
||||
en_bert,
|
||||
style_vec,
|
||||
) in enumerate(train_loader):
|
||||
if net_g.module.use_noise_scaled_mas:
|
||||
if net_g.module.use_noise_scaled_mas and not hps.train.matcha_only:
|
||||
current_mas_noise_scale = (
|
||||
net_g.module.mas_noise_scale_initial
|
||||
- net_g.module.noise_scale_delta * global_step
|
||||
@@ -684,6 +696,63 @@ def train_and_evaluate(
|
||||
en_bert,
|
||||
style_vec,
|
||||
)
|
||||
if hps.train.matcha_only:
|
||||
loss_matcha = matcha_loss.float() * hps.train.c_matcha
|
||||
|
||||
if hps.train.matcha_only:
|
||||
optim_g.zero_grad()
|
||||
scaler.scale(loss_matcha).backward()
|
||||
scaler.unscale_(optim_g)
|
||||
grad_norm_g = commons.clip_grad_value_(
|
||||
net_g.module.matcha.parameters(), None
|
||||
)
|
||||
scaler.step(optim_g)
|
||||
scaler.update()
|
||||
|
||||
if rank == 0:
|
||||
if global_step % hps.train.log_interval == 0 and not hps.speedup:
|
||||
utils.summarize(
|
||||
writer=writer,
|
||||
global_step=global_step,
|
||||
scalars={
|
||||
"loss/g/total": loss_matcha,
|
||||
"loss/g/matcha": loss_matcha,
|
||||
"learning_rate": optim_g.param_groups[0]["lr"],
|
||||
"grad_norm_g": grad_norm_g,
|
||||
},
|
||||
)
|
||||
if (
|
||||
global_step % hps.train.eval_interval == 0
|
||||
and global_step != 0
|
||||
and initial_step != global_step
|
||||
):
|
||||
assert hps.model_dir is not None
|
||||
utils.checkpoints.save_checkpoint(
|
||||
net_g,
|
||||
optim_g,
|
||||
hps.train.learning_rate,
|
||||
epoch,
|
||||
os.path.join(hps.model_dir, f"G_{global_step}.pth"),
|
||||
)
|
||||
utils.safetensors.save_safetensors(
|
||||
net_g,
|
||||
epoch,
|
||||
os.path.join(
|
||||
config.out_dir,
|
||||
f"{config.model_name}_e{epoch}_s{global_step}.safetensors",
|
||||
),
|
||||
for_infer=True,
|
||||
)
|
||||
|
||||
global_step += 1
|
||||
if pbar is not None:
|
||||
pbar.set_description(
|
||||
f"Matcha FT {epoch}({100.0 * batch_idx / len(train_loader):.0f}%)/{hps.train.epochs}"
|
||||
)
|
||||
pbar.update()
|
||||
continue
|
||||
|
||||
with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
|
||||
mel = spec_to_mel_torch(
|
||||
spec,
|
||||
hps.data.filter_length,
|
||||
|
||||
Reference in New Issue
Block a user