# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import torch
import os
import json5
from tqdm import tqdm
import json
import shutil
from models.svc.base import SVCTrainer
from modules.encoder.condition_encoder import ConditionEncoder
from models.svc.comosvc.comosvc import ComoSVC
class ComoSVCTrainer(SVCTrainer):
r"""The base trainer for all diffusion models. It inherits from SVCTrainer and
implements ``_build_model`` and ``_forward_step`` methods.
"""
def __init__(self, args=None, cfg=None):
SVCTrainer.__init__(self, args, cfg)
self.distill = cfg.model.comosvc.distill
self.skip_diff = True
### Following are methods only for comoSVC models ###
def _load_teacher_model(self, model):
r"""Load teacher model from checkpoint file."""
self.checkpoint_file = self.teacher_model_path
self.logger.info(
"Load teacher acoustic model from {}".format(self.checkpoint_file)
)
raw_dict = torch.load(self.checkpoint_file)
model.load_state_dict(raw_dict)
def _build_model(self):
r"""Build the model for training. This function is called in ``__init__`` function."""
# TODO: sort out the config
self.cfg.model.condition_encoder.f0_min = self.cfg.preprocess.f0_min
self.cfg.model.condition_encoder.f0_max = self.cfg.preprocess.f0_max
self.condition_encoder = ConditionEncoder(self.cfg.model.condition_encoder)
self.acoustic_mapper = ComoSVC(self.cfg)
model = torch.nn.ModuleList([self.condition_encoder, self.acoustic_mapper])
if self.cfg.model.comosvc.distill:
if not self.args.resume:
# do not load teacher model when resume
self.teacher_model_path = self.cfg.model.teacher_model_path
self._load_teacher_model(model)
# build teacher & target decoder and freeze teacher
self.acoustic_mapper.decoder.init_consistency_training()
self.freeze_net(self.condition_encoder)
self.freeze_net(self.acoustic_mapper.encoder)
self.freeze_net(self.acoustic_mapper.decoder.denoise_fn_pretrained)
self.freeze_net(self.acoustic_mapper.decoder.denoise_fn_ema)
return model
def freeze_net(self, model):
r"""Freeze the model for training."""
for name, param in model.named_parameters():
param.requires_grad = False
def __build_optimizer(self):
r"""Build optimizer for training. This function is called in ``__init__`` function."""
if self.cfg.train.optimizer.lower() == "adamw":
optimizer = torch.optim.AdamW(
params=filter(lambda p: p.requires_grad, self.model.parameters()),
**self.cfg.train.adamw,
)
else:
raise NotImplementedError(
"Not support optimizer: {}".format(self.cfg.train.optimizer)
)
return optimizer
def _forward_step(self, batch):
r"""Forward step for training and inference. This function is called
in ``_train_step`` & ``_test_step`` function.
"""
loss = {}
mask = batch["mask"]
mel_input = batch["mel"]
cond = self.condition_encoder(batch)
if self.distill:
cond = cond.detach()
self.skip_diff = True if self.step < self.cfg.train.fast_steps else False
ssim_loss, prior_loss, diff_loss = self.acoustic_mapper.compute_loss(
mask, cond, mel_input, skip_diff=self.skip_diff
)
if self.distill:
loss["distil_loss"] = diff_loss
else:
loss["ssim_loss_encoder"] = ssim_loss
loss["prior_loss_encoder"] = prior_loss
loss["diffusion_loss_decoder"] = diff_loss
return loss
def _train_epoch(self):
r"""Training epoch. Should return average loss of a batch (sample) over
one epoch. See ``train_loop`` for usage.
"""
self.model.train()
epoch_sum_loss: float = 0.0
epoch_step: int = 0
for batch in tqdm(
self.train_dataloader,
desc=f"Training Epoch {self.epoch}",
unit="batch",
colour="GREEN",
leave=False,
dynamic_ncols=True,
smoothing=0.04,
disable=not self.accelerator.is_main_process,
):
# Do training step and BP
with self.accelerator.accumulate(self.model):
loss = self._train_step(batch)
total_loss = 0
for k, v in loss.items():
total_loss += v
self.accelerator.backward(total_loss)
enc_grad_norm = torch.nn.utils.clip_grad_norm_(
self.acoustic_mapper.encoder.parameters(), max_norm=1
)
dec_grad_norm = torch.nn.utils.clip_grad_norm_(
self.acoustic_mapper.decoder.parameters(), max_norm=1
)
self.optimizer.step()
self.optimizer.zero_grad()
self.batch_count += 1
# Update info for each step
# TODO: step means BP counts or batch counts?
if self.batch_count % self.cfg.train.gradient_accumulation_step == 0:
epoch_sum_loss += total_loss
log_info = {}
for k, v in loss.items():
key = "Step/Train Loss/{}".format(k)
log_info[key] = v
log_info["Step/Learning Rate"] = self.optimizer.param_groups[0]["lr"]
self.accelerator.log(
log_info,
step=self.step,
)
self.step += 1
epoch_step += 1
self.accelerator.wait_for_everyone()
return (
epoch_sum_loss
/ len(self.train_dataloader)
* self.cfg.train.gradient_accumulation_step,
loss,
)
def train_loop(self):
r"""Training loop. The public entry of training process."""
# Wait everyone to prepare before we move on
self.accelerator.wait_for_everyone()
# dump config file
if self.accelerator.is_main_process:
self.__dump_cfg(self.config_save_path)
self.model.train()
self.optimizer.zero_grad()
# Wait to ensure good to go
self.accelerator.wait_for_everyone()
while self.epoch < self.max_epoch:
self.logger.info("\n")
self.logger.info("-" * 32)
self.logger.info("Epoch {}: ".format(self.epoch))
### TODO: change the return values of _train_epoch() to a loss dict, or (total_loss, loss_dict)
### It's inconvenient for the model with multiple losses
# Do training & validating epoch
train_loss, loss = self._train_epoch()
self.logger.info(" |- Train/Loss: {:.6f}".format(train_loss))
for k, v in loss.items():
self.logger.info(" |- Train/Loss/{}: {:.6f}".format(k, v))
valid_loss = self._valid_epoch()
self.logger.info(" |- Valid/Loss: {:.6f}".format(valid_loss))
self.accelerator.log(
{"Epoch/Train Loss": train_loss, "Epoch/Valid Loss": valid_loss},
step=self.epoch,
)
self.accelerator.wait_for_everyone()
# TODO: what is scheduler?
self.scheduler.step(valid_loss) # FIXME: use epoch track correct?
# Check if hit save_checkpoint_stride and run_eval
run_eval = False
if self.accelerator.is_main_process:
save_checkpoint = False
hit_dix = []
for i, num in enumerate(self.save_checkpoint_stride):
if self.epoch % num == 0:
save_checkpoint = True
hit_dix.append(i)
run_eval |= self.run_eval[i]
self.accelerator.wait_for_everyone()
if (
self.accelerator.is_main_process
and save_checkpoint
and (self.distill or not self.skip_diff)
):
path = os.path.join(
self.checkpoint_dir,
"epoch-{:04d}_step-{:07d}_loss-{:.6f}".format(
self.epoch, self.step, train_loss
),
)
self.tmp_checkpoint_save_path = path
self.accelerator.save_state(path)
print(f"save checkpoint in {path}")
json.dump(
self.checkpoints_path,
open(os.path.join(path, "ckpts.json"), "w"),
ensure_ascii=False,
indent=4,
)
self._save_auxiliary_states()
# Remove old checkpoints
to_remove = []
for idx in hit_dix:
self.checkpoints_path[idx].append(path)
while len(self.checkpoints_path[idx]) > self.keep_last[idx]:
to_remove.append((idx, self.checkpoints_path[idx].pop(0)))
# Search conflicts
total = set()
for i in self.checkpoints_path:
total |= set(i)
do_remove = set()
for idx, path in to_remove[::-1]:
if path in total:
self.checkpoints_path[idx].insert(0, path)
else:
do_remove.add(path)
# Remove old checkpoints
for path in do_remove:
shutil.rmtree(path, ignore_errors=True)
self.logger.debug(f"Remove old checkpoint: {path}")
self.accelerator.wait_for_everyone()
if run_eval:
# TODO: run evaluation
pass
# Update info for each epoch
self.epoch += 1
# Finish training and save final checkpoint
self.accelerator.wait_for_everyone()
if self.accelerator.is_main_process:
self.accelerator.save_state(
os.path.join(
self.checkpoint_dir,
"final_epoch-{:04d}_step-{:07d}_loss-{:.6f}".format(
self.epoch, self.step, valid_loss
),
)
)
self._save_auxiliary_states()
self.accelerator.end_training()
@torch.inference_mode()
def _valid_epoch(self):
r"""Testing epoch. Should return average loss of a batch (sample) over
one epoch. See ``train_loop`` for usage.
"""
self.model.eval()
epoch_sum_loss = 0.0
for batch in tqdm(
self.valid_dataloader,
desc=f"Validating Epoch {self.epoch}",
unit="batch",
colour="GREEN",
leave=False,
dynamic_ncols=True,
smoothing=0.04,
disable=not self.accelerator.is_main_process,
):
batch_loss = self._valid_step(batch)
for k, v in batch_loss.items():
epoch_sum_loss += v
self.accelerator.wait_for_everyone()
return epoch_sum_loss / len(self.valid_dataloader)
@staticmethod
def __count_parameters(model):
model_param = 0.0
if isinstance(model, dict):
for key, value in model.items():
model_param += sum(p.numel() for p in model[key].parameters())
else:
model_param = sum(p.numel() for p in model.parameters())
return model_param
def __dump_cfg(self, path):
os.makedirs(os.path.dirname(path), exist_ok=True)
json5.dump(
self.cfg,
open(path, "w"),
indent=4,
sort_keys=True,
ensure_ascii=False,
quote_keys=True,
)
