svc_trainer.py
# 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 json
import os
import torch
import torch.nn as nn
import numpy as np
from models.base.new_trainer import BaseTrainer
from models.svc.base.svc_dataset import (
SVCOfflineCollator,
SVCOfflineDataset,
SVCOnlineCollator,
SVCOnlineDataset,
)
from processors.audio_features_extractor import AudioFeaturesExtractor
from processors.acoustic_extractor import cal_normalized_mel, load_mel_extrema
EPS = 1.0e-12
class SVCTrainer(BaseTrainer):
r"""The base trainer for all SVC models. It inherits from BaseTrainer and implements
``build_criterion``, ``_build_dataset`` and ``_build_singer_lut`` methods. You can inherit from this
class, and implement ``_build_model``, ``_forward_step``.
"""
def __init__(self, args=None, cfg=None):
self.args = args
self.cfg = cfg
self._init_accelerator()
# Only for SVC tasks
with self.accelerator.main_process_first():
self.singers = self._build_singer_lut()
# Super init
BaseTrainer.__init__(self, args, cfg)
# Only for SVC tasks
self.task_type = "SVC"
self.logger.info("Task type: {}".format(self.task_type))
### Following are methods only for SVC tasks ###
def _build_dataset(self):
self.online_features_extraction = (
self.cfg.preprocess.features_extraction_mode == "online"
)
if not self.online_features_extraction:
return SVCOfflineDataset, SVCOfflineCollator
else:
self.audio_features_extractor = AudioFeaturesExtractor(self.cfg)
return SVCOnlineDataset, SVCOnlineCollator
def _extract_svc_features(self, batch):
"""
Features extraction during training
Batch:
wav: (B, T)
wav_len: (B)
target_len: (B)
mask: (B, n_frames, 1)
spk_id: (B, 1)
wav_{sr}: (B, T)
wav_{sr}_len: (B)
Added elements when output:
mel: (B, n_frames, n_mels)
frame_pitch: (B, n_frames)
frame_uv: (B, n_frames)
frame_energy: (B, n_frames)
frame_{content}: (B, n_frames, D)
"""
padded_n_frames = torch.max(batch["target_len"])
final_n_frames = padded_n_frames
### Mel Spectrogram ###
if self.cfg.preprocess.use_mel:
# (B, n_mels, n_frames)
raw_mel = self.audio_features_extractor.get_mel_spectrogram(batch["wav"])
if self.cfg.preprocess.use_min_max_norm_mel:
# TODO: Change the hard code
# Using the empirical mel extrema to denormalize
if not hasattr(self, "mel_extrema"):
# (n_mels)
m, M = load_mel_extrema(self.cfg.preprocess, "vctk")
# (1, n_mels, 1)
m = (
torch.as_tensor(m, device=raw_mel.device)
.unsqueeze(0)
.unsqueeze(-1)
)
M = (
torch.as_tensor(M, device=raw_mel.device)
.unsqueeze(0)
.unsqueeze(-1)
)
self.mel_extrema = m, M
m, M = self.mel_extrema
mel = (raw_mel - m) / (M - m + EPS) * 2 - 1
else:
mel = raw_mel
final_n_frames = min(final_n_frames, mel.size(-1))
# (B, n_frames, n_mels)
batch["mel"] = mel.transpose(1, 2)
else:
raw_mel = None
### F0 ###
if self.cfg.preprocess.use_frame_pitch:
# (B, n_frames)
raw_f0, raw_uv = self.audio_features_extractor.get_f0(
batch["wav"],
wav_lens=batch["wav_len"],
use_interpolate=self.cfg.preprocess.use_interpolation_for_uv,
return_uv=True,
)
final_n_frames = min(final_n_frames, raw_f0.size(-1))
batch["frame_pitch"] = raw_f0
if self.cfg.preprocess.use_uv:
batch["frame_uv"] = raw_uv
### Energy ###
if self.cfg.preprocess.use_frame_energy:
# (B, n_frames)
raw_energy = self.audio_features_extractor.get_energy(
batch["wav"], mel_spec=raw_mel
)
final_n_frames = min(final_n_frames, raw_energy.size(-1))
batch["frame_energy"] = raw_energy
### Semantic Features ###
if self.cfg.model.condition_encoder.use_whisper:
# (B, n_frames, D)
whisper_feats = self.audio_features_extractor.get_whisper_features(
wavs=batch["wav_{}".format(self.cfg.preprocess.whisper_sample_rate)],
target_frame_len=padded_n_frames,
)
final_n_frames = min(final_n_frames, whisper_feats.size(1))
batch["whisper_feat"] = whisper_feats
if self.cfg.model.condition_encoder.use_contentvec:
# (B, n_frames, D)
contentvec_feats = self.audio_features_extractor.get_contentvec_features(
wavs=batch["wav_{}".format(self.cfg.preprocess.contentvec_sample_rate)],
target_frame_len=padded_n_frames,
)
final_n_frames = min(final_n_frames, contentvec_feats.size(1))
batch["contentvec_feat"] = contentvec_feats
if self.cfg.model.condition_encoder.use_wenet:
# (B, n_frames, D)
wenet_feats = self.audio_features_extractor.get_wenet_features(
wavs=batch["wav_{}".format(self.cfg.preprocess.wenet_sample_rate)],
target_frame_len=padded_n_frames,
wav_lens=batch[
"wav_{}_len".format(self.cfg.preprocess.wenet_sample_rate)
],
)
final_n_frames = min(final_n_frames, wenet_feats.size(1))
batch["wenet_feat"] = wenet_feats
### Align all the audio features to the same frame length ###
frame_level_features = [
"mask",
"mel",
"frame_pitch",
"frame_uv",
"frame_energy",
"whisper_feat",
"contentvec_feat",
"wenet_feat",
]
for k in frame_level_features:
if k in batch:
# (B, n_frames, ...)
batch[k] = batch[k][:, :final_n_frames].contiguous()
return batch
@staticmethod
def _build_criterion():
criterion = nn.MSELoss(reduction="none")
return criterion
@staticmethod
def _compute_loss(criterion, y_pred, y_gt, loss_mask):
"""
Args:
criterion: MSELoss(reduction='none')
y_pred, y_gt: (B, seq_len, D)
loss_mask: (B, seq_len, 1)
Returns:
loss: Tensor of shape []
"""
# (B, seq_len, D)
loss = criterion(y_pred, y_gt)
# expand loss_mask to (B, seq_len, D)
loss_mask = loss_mask.repeat(1, 1, loss.shape[-1])
loss = torch.sum(loss * loss_mask) / torch.sum(loss_mask)
return loss
def _save_auxiliary_states(self):
"""
To save the singer's look-up table in the checkpoint saving path
"""
with open(
os.path.join(self.tmp_checkpoint_save_path, self.cfg.preprocess.spk2id),
"w",
encoding="utf-8",
) as f:
json.dump(self.singers, f, indent=4, ensure_ascii=False)
def _build_singer_lut(self):
resumed_singer_path = None
if self.args.resume_from_ckpt_path and self.args.resume_from_ckpt_path != "":
resumed_singer_path = os.path.join(
self.args.resume_from_ckpt_path, self.cfg.preprocess.spk2id
)
if os.path.exists(os.path.join(self.exp_dir, self.cfg.preprocess.spk2id)):
resumed_singer_path = os.path.join(self.exp_dir, self.cfg.preprocess.spk2id)
if resumed_singer_path:
with open(resumed_singer_path, "r") as f:
singers = json.load(f)
else:
singers = dict()
for dataset in self.cfg.dataset:
singer_lut_path = os.path.join(
self.cfg.preprocess.processed_dir, dataset, self.cfg.preprocess.spk2id
)
with open(singer_lut_path, "r") as singer_lut_path:
singer_lut = json.load(singer_lut_path)
for singer in singer_lut.keys():
if singer not in singers:
singers[singer] = len(singers)
with open(
os.path.join(self.exp_dir, self.cfg.preprocess.spk2id), "w"
) as singer_file:
json.dump(singers, singer_file, indent=4, ensure_ascii=False)
print(
"singers have been dumped to {}".format(
os.path.join(self.exp_dir, self.cfg.preprocess.spk2id)
)
)
return singers
