import os import sys import glob import json import torch import hashlib import logging import argparse import datetime import warnings import torch.distributed as dist import torch.utils.data as tdata import torch.multiprocessing as mp from tqdm import tqdm from collections import OrderedDict from random import randint, shuffle from torch.amp import GradScaler, autocast from torch.utils.tensorboard import SummaryWriter from time import time as ttime from torch.nn import functional as F from distutils.util import strtobool from torch.nn.parallel import DistributedDataParallel as DDP sys.path.append(os.getcwd()) from main.library import opencl from main.app.variables import logger, translations from main.inference.conversion.utils import clear_gpu_cache from main.library.algorithm.synthesizers import Synthesizer from main.library.algorithm.discriminators import MultiPeriodDiscriminator from main.library.algorithm.commons import slice_segments, clip_grad_value from main.inference.training.mel_processing import spec_to_mel_torch, mel_spectrogram_torch from main.inference.training.losses import discriminator_loss, kl_loss, feature_loss, generator_loss from main.inference.training.data_utils import TextAudioCollate, TextAudioCollateMultiNSFsid, TextAudioLoader, TextAudioLoaderMultiNSFsid, DistributedBucketSampler from main.inference.training.utils import HParams, replace_keys_in_dict, load_checkpoint, latest_checkpoint_path, save_checkpoint, summarize, plot_spectrogram_to_numpy from main.app.variables import config as main_config from main.app.variables import configs as main_configs warnings.filterwarnings("ignore") logging.getLogger("torch").setLevel(logging.ERROR) def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument("--train", action='store_true') parser.add_argument("--model_name", type=str, required=True) parser.add_argument("--rvc_version", type=str, default="v2") parser.add_argument("--save_every_epoch", type=int, required=True) parser.add_argument("--save_only_latest", type=lambda x: bool(strtobool(x)), default=True) parser.add_argument("--save_every_weights", type=lambda x: bool(strtobool(x)), default=True) parser.add_argument("--total_epoch", type=int, default=300) parser.add_argument("--sample_rate", type=int, required=True) parser.add_argument("--batch_size", type=int, default=8) parser.add_argument("--gpu", type=str, default="0") parser.add_argument("--pitch_guidance", type=lambda x: bool(strtobool(x)), default=True) parser.add_argument("--g_pretrained_path", type=str, default="") parser.add_argument("--d_pretrained_path", type=str, default="") parser.add_argument("--overtraining_detector", type=lambda x: bool(strtobool(x)), default=False) parser.add_argument("--overtraining_threshold", type=int, default=50) parser.add_argument("--cleanup", type=lambda x: bool(strtobool(x)), default=False) parser.add_argument("--cache_data_in_gpu", type=lambda x: bool(strtobool(x)), default=False) parser.add_argument("--model_author", type=str) parser.add_argument("--vocoder", type=str, default="Default") parser.add_argument("--checkpointing", type=lambda x: bool(strtobool(x)), default=False) parser.add_argument("--deterministic", type=lambda x: bool(strtobool(x)), default=False) parser.add_argument("--benchmark", type=lambda x: bool(strtobool(x)), default=False) parser.add_argument("--optimizer", type=str, default="AdamW") parser.add_argument("--energy_use", type=lambda x: bool(strtobool(x)), default=False) return parser.parse_args() args = parse_arguments() model_name, save_every_epoch, total_epoch, pretrainG, pretrainD, version, gpus, batch_size, sample_rate, pitch_guidance, save_only_latest, save_every_weights, cache_data_in_gpu, overtraining_detector, overtraining_threshold, cleanup, model_author, vocoder, checkpointing, optimizer_choice, energy_use = args.model_name, args.save_every_epoch, args.total_epoch, args.g_pretrained_path, args.d_pretrained_path, args.rvc_version, args.gpu, args.batch_size, args.sample_rate, args.pitch_guidance, args.save_only_latest, args.save_every_weights, args.cache_data_in_gpu, args.overtraining_detector, args.overtraining_threshold, args.cleanup, args.model_author, args.vocoder, args.checkpointing, args.optimizer, args.energy_use experiment_dir = os.path.join(main_configs["logs_path"], model_name) training_file_path = os.path.join(experiment_dir, "training_data.json") config_save_path = os.path.join(experiment_dir, "config.json") torch.backends.cudnn.deterministic = args.deterministic if not main_config.device.startswith("ocl") else False torch.backends.cudnn.benchmark = args.benchmark if not main_config.device.startswith("ocl") else False lowest_value = {"step": 0, "value": float("inf"), "epoch": 0} global_step, last_loss_gen_all, overtrain_save_epoch = 0, 0, 0 loss_gen_history, smoothed_loss_gen_history, loss_disc_history, smoothed_loss_disc_history = [], [], [], [] with open(config_save_path, "r") as f: config = json.load(f) config = HParams(**config) config.data.training_files = os.path.join(experiment_dir, "filelist.txt") def main(): global training_file_path, last_loss_gen_all, smoothed_loss_gen_history, loss_gen_history, loss_disc_history, smoothed_loss_disc_history, overtrain_save_epoch, model_author, vocoder, checkpointing, gpus, energy_use log_data = {translations['modelname']: model_name, translations["save_every_epoch"]: save_every_epoch, translations["total_e"]: total_epoch, translations["dorg"].format(pretrainG=pretrainG, pretrainD=pretrainD): "", translations['training_version']: version, "Gpu": gpus, translations['batch_size']: batch_size, translations['pretrain_sr']: sample_rate, translations['training_f0']: pitch_guidance, translations['save_only_latest']: save_only_latest, translations['save_every_weights']: save_every_weights, translations['cache_in_gpu']: cache_data_in_gpu, translations['overtraining_detector']: overtraining_detector, translations['threshold']: overtraining_threshold, translations['cleanup_training']: cleanup, translations['memory_efficient_training']: checkpointing, translations["optimizer"]: optimizer_choice, translations["train&energy"]: energy_use} if model_author: log_data[translations["model_author"].format(model_author=model_author)] = "" if vocoder != "Default": log_data[translations['vocoder']] = vocoder for key, value in log_data.items(): logger.debug(f"{key}: {value}" if value != "" else f"{key} {value}") try: os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(randint(20000, 55555)) if torch.cuda.is_available(): device, gpus = torch.device("cuda"), [int(item) for item in gpus.split("-")] n_gpus = len(gpus) elif opencl.is_available(): device, gpus = torch.device("ocl"), [int(item) for item in gpus.split("-")] n_gpus = len(gpus) elif torch.backends.mps.is_available(): device, gpus = torch.device("mps"), [0] n_gpus = 1 else: device, gpus = torch.device("cpu"), [0] n_gpus = 1 logger.warning(translations["not_gpu"]) def start(): children = [] pid_data = {"process_pids": []} with open(config_save_path, "r") as pid_file: try: pid_data.update(json.load(pid_file)) except json.JSONDecodeError: pass with open(config_save_path, "w") as pid_file: for rank, device_id in enumerate(gpus): subproc = mp.Process(target=run, args=(rank, n_gpus, experiment_dir, pretrainG, pretrainD, pitch_guidance, total_epoch, save_every_weights, config, device, device_id, model_author, vocoder, checkpointing, energy_use)) children.append(subproc) subproc.start() pid_data["process_pids"].append(subproc.pid) json.dump(pid_data, pid_file, indent=4) for i in range(n_gpus): children[i].join() def load_from_json(file_path): if os.path.exists(file_path): with open(file_path, "r") as f: data = json.load(f) return (data.get("loss_disc_history", []), data.get("smoothed_loss_disc_history", []), data.get("loss_gen_history", []), data.get("smoothed_loss_gen_history", [])) return [], [], [], [] def continue_overtrain_detector(training_file_path): if overtraining_detector and os.path.exists(training_file_path): (loss_disc_history, smoothed_loss_disc_history, loss_gen_history, smoothed_loss_gen_history) = load_from_json(training_file_path) if cleanup: for root, dirs, files in os.walk(experiment_dir, topdown=False): for name in files: file_path = os.path.join(root, name) _, file_extension = os.path.splitext(name) if (file_extension == ".0" or (name.startswith("D_") and file_extension == ".pth") or (name.startswith("G_") and file_extension == ".pth") or (file_extension == ".index")): os.remove(file_path) for name in dirs: if name == "eval": folder_path = os.path.join(root, name) for item in os.listdir(folder_path): item_path = os.path.join(folder_path, item) if os.path.isfile(item_path): os.remove(item_path) os.rmdir(folder_path) continue_overtrain_detector(training_file_path) start() except Exception as e: logger.error(f"{translations['training_error']} {e}") import traceback logger.debug(traceback.format_exc()) def verify_checkpoint_shapes(checkpoint_path, model): checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) checkpoint_state_dict = checkpoint["model"] try: model_state_dict = model.module.load_state_dict(checkpoint_state_dict) if hasattr(model, "module") else model.load_state_dict(checkpoint_state_dict) except RuntimeError: logger.warning(translations["checkpointing_err"]) sys.exit(1) else: del checkpoint, checkpoint_state_dict, model_state_dict class EpochRecorder: def __init__(self): self.last_time = ttime() def record(self): now_time = ttime() elapsed_time = now_time - self.last_time self.last_time = now_time return translations["time_or_speed_training"].format(current_time=datetime.datetime.now().strftime("%H:%M:%S"), elapsed_time_str=str(datetime.timedelta(seconds=int(round(elapsed_time, 1))))) def extract_model(ckpt, sr, pitch_guidance, name, model_path, epoch, step, version, hps, model_author, vocoder, energy_use): try: logger.info(translations["savemodel"].format(model_dir=model_path, epoch=epoch, step=step)) os.makedirs(os.path.dirname(model_path), exist_ok=True) opt = OrderedDict(weight={key: value.half() for key, value in ckpt.items() if "enc_q" not in key}) opt["config"] = [hps.data.filter_length // 2 + 1, 32, hps.model.inter_channels, hps.model.hidden_channels, hps.model.filter_channels, hps.model.n_heads, hps.model.n_layers, hps.model.kernel_size, hps.model.p_dropout, hps.model.resblock, hps.model.resblock_kernel_sizes, hps.model.resblock_dilation_sizes, hps.model.upsample_rates, hps.model.upsample_initial_channel, hps.model.upsample_kernel_sizes, hps.model.spk_embed_dim, hps.model.gin_channels, hps.data.sample_rate] opt["epoch"] = f"{epoch}epoch" opt["step"] = step opt["sr"] = sr opt["f0"] = int(pitch_guidance) opt["version"] = version opt["creation_date"] = datetime.datetime.now().isoformat() opt["model_hash"] = hashlib.sha256(f"{str(ckpt)} {epoch} {step} {datetime.datetime.now().isoformat()}".encode()).hexdigest() opt["model_name"] = name opt["author"] = model_author opt["vocoder"] = vocoder opt["energy"] = energy_use torch.save(replace_keys_in_dict(replace_keys_in_dict(opt, ".parametrizations.weight.original1", ".weight_v"), ".parametrizations.weight.original0", ".weight_g"), model_path) except Exception as e: logger.error(f"{translations['extract_model_error']}: {e}") def run(rank, n_gpus, experiment_dir, pretrainG, pretrainD, pitch_guidance, custom_total_epoch, custom_save_every_weights, config, device, device_id, model_author, vocoder, checkpointing, energy_use): global global_step, optimizer_choice try: dist.init_process_group(backend=("gloo" if sys.platform == "win32" or device.type != "cuda" else "nccl"), init_method="env://", world_size=n_gpus, rank=rank) except: dist.init_process_group(backend=("gloo" if sys.platform == "win32" or device.type != "cuda" else "nccl"), init_method="env://?use_libuv=False", world_size=n_gpus, rank=rank) torch.manual_seed(config.train.seed) if device.type == "cuda": torch.cuda.manual_seed(config.train.seed) elif device.type == "ocl": opencl.pytorch_ocl.manual_seed_all(config.train.seed) if torch.cuda.is_available(): torch.cuda.set_device(device_id) writer_eval = SummaryWriter(log_dir=os.path.join(experiment_dir, "eval")) if rank == 0 else None if pitch_guidance: train_dataset = TextAudioLoaderMultiNSFsid(config.data, energy=energy_use) collate_fn = TextAudioCollateMultiNSFsid(energy=energy_use) else: train_dataset = TextAudioLoader(config.data, energy=energy_use) collate_fn = TextAudioCollate(energy=energy_use) train_loader = tdata.DataLoader(train_dataset, num_workers=4, shuffle=False, pin_memory=True, collate_fn=collate_fn, batch_sampler=DistributedBucketSampler(train_dataset, batch_size * n_gpus, [100, 200, 300, 400, 500, 600, 700, 800, 900], num_replicas=n_gpus, rank=rank, shuffle=True), persistent_workers=True, prefetch_factor=8) net_g, net_d = Synthesizer(config.data.filter_length // 2 + 1, config.train.segment_size // config.data.hop_length, **config.model, use_f0=pitch_guidance, sr=sample_rate, vocoder=vocoder, checkpointing=checkpointing, energy=energy_use), MultiPeriodDiscriminator(version, config.model.use_spectral_norm, checkpointing=checkpointing) net_g, net_d = (net_g.cuda(device_id), net_d.cuda(device_id)) if torch.cuda.is_available() else (net_g.to(device), net_d.to(device)) optimizer_optim = torch.optim.AdamW if optimizer_choice == "AdamW" else torch.optim.RAdam optim_g, optim_d = optimizer_optim(net_g.parameters(), config.train.learning_rate, betas=config.train.betas, eps=config.train.eps), optimizer_optim(net_d.parameters(), config.train.learning_rate, betas=config.train.betas, eps=config.train.eps) if device.type != "ocl": net_g, net_d = (DDP(net_g, device_ids=[device_id]), DDP(net_d, device_ids=[device_id])) if torch.cuda.is_available() else (DDP(net_g), DDP(net_d)) try: logger.info(translations["start_training"]) _, _, _, epoch_str = load_checkpoint(logger, (os.path.join(experiment_dir, "D_latest.pth") if save_only_latest else latest_checkpoint_path(experiment_dir, "D_*.pth")), net_d, optim_d) _, _, _, epoch_str = load_checkpoint(logger, (os.path.join(experiment_dir, "G_latest.pth") if save_only_latest else latest_checkpoint_path(experiment_dir, "G_*.pth")), net_g, optim_g) epoch_str += 1 global_step = (epoch_str - 1) * len(train_loader) except: epoch_str, global_step = 1, 0 verify = main_configs.get("pretrain_verify_shape", True) strict = main_configs.get("pretrain_strict", True) if pretrainG != "" and pretrainG != "None": if rank == 0: if verify: verify_checkpoint_shapes(pretrainG, net_g) logger.info(translations["import_pretrain"].format(dg="G", pretrain=pretrainG)) ckptG = torch.load(pretrainG, map_location="cpu", weights_only=True)["model"] net_g.module.load_state_dict(ckptG, strict=strict) if hasattr(net_g, "module") else net_g.load_state_dict(ckptG, strict=strict) else: logger.warning(translations["not_using_pretrain"].format(dg="G")) if pretrainD != "" and pretrainD != "None": if rank == 0: if verify: verify_checkpoint_shapes(pretrainD, net_d) logger.info(translations["import_pretrain"].format(dg="D", pretrain=pretrainD)) ckptD = torch.load(pretrainD, map_location="cpu", weights_only=True)["model"] net_d.module.load_state_dict(ckptD, strict=strict) if hasattr(net_d, "module") else net_d.load_state_dict(ckptD, strict=strict) else: logger.warning(translations["not_using_pretrain"].format(dg="D")) scheduler_g, scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=config.train.lr_decay, last_epoch=epoch_str - 2), torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=config.train.lr_decay, last_epoch=epoch_str - 2) scaler = GradScaler(device=device, enabled=main_config.is_half and device.type == "cuda") optim_g.step(); optim_d.step() cache = [] def to_device(x): return x.cuda(device_id, non_blocking=True) if device.type == "cuda" else x.to(device) for info in train_loader: reference = (to_device(info[0]), to_device(info[1])) if pitch_guidance: reference += (to_device(info[2]), to_device(info[3]), to_device(info[8])) reference += (to_device(info[9]),) if energy_use else (None,) else: reference += (None, None, to_device(info[6])) reference += (to_device(info[7]),) if energy_use else (None,) break for epoch in range(epoch_str, total_epoch + 1): train_and_evaluate(rank, epoch, config, [net_g, net_d], [optim_g, optim_d], scaler, train_loader, writer_eval, cache, custom_save_every_weights, custom_total_epoch, device, device_id, reference, model_author, vocoder, energy_use) scheduler_g.step(); scheduler_d.step() def train_and_evaluate(rank, epoch, hps, nets, optims, scaler, train_loader, writer, cache, custom_save_every_weights, custom_total_epoch, device, device_id, reference, model_author, vocoder, energy_use): global global_step, lowest_value, loss_disc, consecutive_increases_gen, consecutive_increases_disc if epoch == 1: lowest_value = {"step": 0, "value": float("inf"), "epoch": 0} last_loss_gen_all, consecutive_increases_gen, consecutive_increases_disc = 0.0, 0, 0 net_g, net_d = nets optim_g, optim_d = optims train_loader.batch_sampler.set_epoch(epoch) net_g.train(); net_d.train() if device.type == "cuda" and cache_data_in_gpu: data_iterator = cache if cache == []: for batch_idx, info in enumerate(train_loader): cache.append((batch_idx, [tensor.cuda(device_id, non_blocking=True) for tensor in info])) else: shuffle(cache) elif device.type == "ocl" and cache_data_in_gpu: data_iterator = cache if cache == []: for batch_idx, info in enumerate(train_loader): cache.append((batch_idx, [tensor.to(device_id, non_blocking=True) for tensor in info])) else: shuffle(cache) else: data_iterator = enumerate(train_loader) epoch_recorder = EpochRecorder() autocast_enabled = main_config.is_half and device.type == "cuda" autocast_device = "cpu" if str(device.type).startswith("ocl") else device.type autocast_dtype = torch.float32 if not autocast_enabled else (torch.bfloat16 if main_config.brain else torch.float16) with tqdm(total=len(train_loader), leave=False) as pbar: for batch_idx, info in data_iterator: if device.type == "cuda" and not cache_data_in_gpu: info = [tensor.cuda(device_id, non_blocking=True) for tensor in info] elif device.type == "ocl" and not cache_data_in_gpu: info = [tensor.to(device_id, non_blocking=True) for tensor in info] else: info = [tensor.to(device) for tensor in info] phone, phone_lengths = info[0], info[1] if pitch_guidance: pitch, pitchf = info[2], info[3] spec, spec_lengths, wave, sid = info[4], info[5], info[6], info[8] energy = info[9] if energy_use else None else: pitch = pitchf = None spec, spec_lengths, wave, sid = info[2], info[3], info[4], info[6] energy = info[7] if energy_use else None with autocast(autocast_device , enabled=autocast_enabled, dtype=autocast_dtype): y_hat, ids_slice, _, z_mask, (_, z_p, m_p, logs_p, _, logs_q) = net_g(phone, phone_lengths, pitch, pitchf, spec, spec_lengths, sid, energy) mel = spec_to_mel_torch(spec, config.data.filter_length, config.data.n_mel_channels, config.data.sample_rate, config.data.mel_fmin, config.data.mel_fmax) y_mel = slice_segments(mel, ids_slice, config.train.segment_size // config.data.hop_length, dim=3) with autocast(autocast_device, enabled=autocast_enabled, dtype=autocast_dtype): y_hat_mel = mel_spectrogram_torch(y_hat.float().squeeze(1), config.data.filter_length, config.data.n_mel_channels, config.data.sample_rate, config.data.hop_length, config.data.win_length, config.data.mel_fmin, config.data.mel_fmax) wave = slice_segments(wave, ids_slice * config.data.hop_length, config.train.segment_size, dim=3) y_d_hat_r, y_d_hat_g, _, _ = net_d(wave, y_hat.detach()) with autocast(autocast_device, enabled=autocast_enabled, dtype=autocast_dtype): loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g) optim_d.zero_grad() scaler.scale(loss_disc).backward() scaler.unscale_(optim_d) grad_norm_d = clip_grad_value(net_d.parameters(), None) scaler.step(optim_d) with autocast(autocast_device, enabled=autocast_enabled, dtype=autocast_dtype): y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(wave, y_hat) with autocast(autocast_device, enabled=autocast_enabled, dtype=autocast_dtype): loss_mel = F.l1_loss(y_mel, y_hat_mel) * config.train.c_mel loss_kl = (kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * config.train.c_kl) loss_fm = feature_loss(fmap_r, fmap_g) loss_gen, losses_gen = generator_loss(y_d_hat_g) loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl if loss_gen_all < lowest_value["value"]: lowest_value = {"step": global_step, "value": loss_gen_all, "epoch": epoch} optim_g.zero_grad() scaler.scale(loss_gen_all).backward() scaler.unscale_(optim_g) grad_norm_g = clip_grad_value(net_g.parameters(), None) scaler.step(optim_g) scaler.update() if rank == 0 and global_step % config.train.log_interval == 0: if loss_mel > 75: loss_mel = 75 if loss_kl > 9: loss_kl = 9 scalar_dict = {"loss/g/total": loss_gen_all, "loss/d/total": loss_disc, "learning_rate": optim_g.param_groups[0]["lr"], "grad/norm_d": grad_norm_d, "grad/norm_g": grad_norm_g, "loss/g/fm": loss_fm, "loss/g/mel": loss_mel, "loss/g/kl": loss_kl} scalar_dict.update({f"loss/g/{i}": v for i, v in enumerate(losses_gen)}) scalar_dict.update({f"loss/d_r/{i}": v for i, v in enumerate(losses_disc_r)}) scalar_dict.update({f"loss/d_g/{i}": v for i, v in enumerate(losses_disc_g)}) with torch.no_grad(): o, *_ = net_g.module.infer(*reference) if hasattr(net_g, "module") else net_g.infer(*reference) summarize(writer=writer, global_step=global_step, images={"slice/mel_org": plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()), "slice/mel_gen": plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()), "all/mel": plot_spectrogram_to_numpy(mel[0].data.cpu().numpy())}, scalars=scalar_dict, audios={f"gen/audio_{global_step:07d}": o[0, :, :]}, audio_sample_rate=config.data.sample_rate) global_step += 1 pbar.update(1) with torch.no_grad(): clear_gpu_cache() def check_overtraining(smoothed_loss_history, threshold, epsilon=0.004): if len(smoothed_loss_history) < threshold + 1: return False for i in range(-threshold, -1): if smoothed_loss_history[i + 1] > smoothed_loss_history[i]: return True if abs(smoothed_loss_history[i + 1] - smoothed_loss_history[i]) >= epsilon: return False return True def update_exponential_moving_average(smoothed_loss_history, new_value, smoothing=0.987): smoothed_value = new_value if not smoothed_loss_history else (smoothing * smoothed_loss_history[-1] + (1 - smoothing) * new_value) smoothed_loss_history.append(smoothed_value) return smoothed_value def save_to_json(file_path, loss_disc_history, smoothed_loss_disc_history, loss_gen_history, smoothed_loss_gen_history): with open(file_path, "w") as f: json.dump({"loss_disc_history": loss_disc_history, "smoothed_loss_disc_history": smoothed_loss_disc_history, "loss_gen_history": loss_gen_history, "smoothed_loss_gen_history": smoothed_loss_gen_history}, f) model_add, model_del = [], [] done = False if rank == 0: if epoch % save_every_epoch == False: checkpoint_suffix = f"{'latest' if save_only_latest else global_step}.pth" save_checkpoint(logger, net_g, optim_g, config.train.learning_rate, epoch, os.path.join(experiment_dir, "G_" + checkpoint_suffix)) save_checkpoint(logger, net_d, optim_d, config.train.learning_rate, epoch, os.path.join(experiment_dir, "D_" + checkpoint_suffix)) if custom_save_every_weights: model_add.append(os.path.join(main_configs["weights_path"], f"{model_name}_{epoch}e_{global_step}s.pth")) if overtraining_detector and epoch > 1: current_loss_disc, current_loss_gen = float(loss_disc), float(lowest_value["value"]) loss_disc_history.append(current_loss_disc) loss_gen_history.append(current_loss_gen) smoothed_value_disc = update_exponential_moving_average(smoothed_loss_disc_history, current_loss_disc) smoothed_value_gen = update_exponential_moving_average(smoothed_loss_gen_history, current_loss_gen) is_overtraining_disc = check_overtraining(smoothed_loss_disc_history, overtraining_threshold * 2) is_overtraining_gen = check_overtraining(smoothed_loss_gen_history, overtraining_threshold, 0.01) consecutive_increases_disc = (consecutive_increases_disc + 1) if is_overtraining_disc else 0 consecutive_increases_gen = (consecutive_increases_gen + 1) if is_overtraining_gen else 0 if epoch % save_every_epoch == 0: save_to_json(training_file_path, loss_disc_history, smoothed_loss_disc_history, loss_gen_history, smoothed_loss_gen_history) if (is_overtraining_gen and consecutive_increases_gen == overtraining_threshold or is_overtraining_disc and consecutive_increases_disc == (overtraining_threshold * 2)): logger.info(translations["overtraining_find"].format(epoch=epoch, smoothed_value_gen=f"{smoothed_value_gen:.3f}", smoothed_value_disc=f"{smoothed_value_disc:.3f}")) done = True else: logger.info(translations["best_epoch"].format(epoch=epoch, smoothed_value_gen=f"{smoothed_value_gen:.3f}", smoothed_value_disc=f"{smoothed_value_disc:.3f}")) for file in glob.glob(os.path.join(main_configs["weights_path"], f"{model_name}_*e_*s_best_epoch.pth")): model_del.append(file) model_add.append(os.path.join(main_configs["weights_path"], f"{model_name}_{epoch}e_{global_step}s_best_epoch.pth")) if epoch >= custom_total_epoch: logger.info(translations["success_training"].format(epoch=epoch, global_step=global_step, loss_gen_all=round(loss_gen_all.item(), 3))) logger.info(translations["training_info"].format(lowest_value_rounded=round(float(lowest_value["value"]), 3), lowest_value_epoch=lowest_value['epoch'], lowest_value_step=lowest_value['step'])) model_add.append(os.path.join(main_configs["weights_path"], f"{model_name}_{epoch}e_{global_step}s.pth")) done = True for m in model_del: os.remove(m) if model_add: ckpt = (net_g.module.state_dict() if hasattr(net_g, "module") else net_g.state_dict()) for m in model_add: extract_model(ckpt=ckpt, sr=sample_rate, pitch_guidance=pitch_guidance == True, name=model_name, model_path=m, epoch=epoch, step=global_step, version=version, hps=hps, model_author=model_author, vocoder=vocoder, energy_use=energy_use) lowest_value_rounded = round(float(lowest_value["value"]), 3) if epoch > 1 and overtraining_detector: logger.info(translations["model_training_info"].format(model_name=model_name, epoch=epoch, global_step=global_step, epoch_recorder=epoch_recorder.record(), lowest_value_rounded=lowest_value_rounded, lowest_value_epoch=lowest_value['epoch'], lowest_value_step=lowest_value['step'], remaining_epochs_gen=(overtraining_threshold - consecutive_increases_gen), remaining_epochs_disc=((overtraining_threshold * 2) - consecutive_increases_disc), smoothed_value_gen=f"{smoothed_value_gen:.3f}", smoothed_value_disc=f"{smoothed_value_disc:.3f}")) elif epoch > 1 and overtraining_detector == False: logger.info(translations["model_training_info_2"].format(model_name=model_name, epoch=epoch, global_step=global_step, epoch_recorder=epoch_recorder.record(), lowest_value_rounded=lowest_value_rounded, lowest_value_epoch=lowest_value['epoch'], lowest_value_step=lowest_value['step'])) else: logger.info(translations["model_training_info_3"].format(model_name=model_name, epoch=epoch, global_step=global_step, epoch_recorder=epoch_recorder.record())) logger.debug(f"loss_gen_all: {loss_gen_all} loss_gen: {loss_gen} loss_fm: {loss_fm} loss_mel: {loss_mel} loss_kl: {loss_kl}") last_loss_gen_all = loss_gen_all if done: pid_file_path = os.path.join(experiment_dir, "config.json") with open(pid_file_path, "r") as pid_file: pid_data = json.load(pid_file) with open(pid_file_path, "w") as pid_file: pid_data.pop("process_pids", None) json.dump(pid_data, pid_file, indent=4) if os.path.exists(os.path.join(experiment_dir, "train_pid.txt")): os.remove(os.path.join(experiment_dir, "train_pid.txt")) os._exit(0) with torch.no_grad(): clear_gpu_cache() if __name__ == "__main__": torch.multiprocessing.set_start_method("spawn") main()