Upload 5 files

gradient_reversal.py

@@ -0,0 +1,35 @@
+# 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.
+
+from torch.autograd import Function
+import torch
+from torch import nn
+
+
+class GradientReversal(Function):
+    @staticmethod
+    def forward(ctx, x, alpha):
+        ctx.save_for_backward(x, alpha)
+        return x
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = None
+        _, alpha = ctx.saved_tensors
+        if ctx.needs_input_grad[0]:
+            grad_input = -alpha * grad_output
+        return grad_input, None
+
+
+revgrad = GradientReversal.apply
+
+
+class GradientReversal(nn.Module):
+    def __init__(self, alpha):
+        super().__init__()
+        self.alpha = torch.tensor(alpha, requires_grad=False)
+
+    def forward(self, x):
+        return revgrad(x, self.alpha)

losses.py

@@ -0,0 +1,309 @@
+import torch
+import torch.nn.functional as F
+from torchaudio.transforms import MelSpectrogram
+
+
+def adversarial_g_loss(y_disc_gen):
+    """Hinge loss"""
+    loss = 0.0
+    for i in range(len(y_disc_gen)):
+        stft_loss = F.relu(1 - y_disc_gen[i]).mean().squeeze()
+        loss += stft_loss
+    return loss / len(y_disc_gen)
+
+
+def feature_loss(fmap_r, fmap_gen):
+    loss = 0.0
+    for i in range(len(fmap_r)):
+        for j in range(len(fmap_r[i])):
+            stft_loss = ((fmap_r[i][j] - fmap_gen[i][j]).abs() /
+                         (fmap_r[i][j].abs().mean())).mean()
+            loss += stft_loss
+    return loss / (len(fmap_r) * len(fmap_r[0]))
+
+
+def sim_loss(y_disc_r, y_disc_gen):
+    loss = 0.0
+    for i in range(len(y_disc_r)):
+        loss += F.mse_loss(y_disc_r[i], y_disc_gen[i])
+    return loss / len(y_disc_r)
+
+# def sisnr_loss(x, s, eps=1e-8):
+    # """
+    # calculate training loss
+    # input:
+          # x: separated signal, N x S tensor, estimate value
+          # s: reference signal, N x S tensor, True value
+    # Return:
+          # sisnr: N tensor
+    # """
+    # if x.shape != s.shape:
+        # if x.shape[-1] > s.shape[-1]:
+            # x = x[:, :s.shape[-1]]
+        # else:
+            # s = s[:, :x.shape[-1]]
+    # def l2norm(mat, keepdim=False):
+        # return torch.norm(mat, dim=-1, keepdim=keepdim)
+    # if x.shape != s.shape:
+        # raise RuntimeError(
+            # "Dimention mismatch when calculate si-snr, {} vs {}".format(
+                # x.shape, s.shape))
+    # x_zm = x - torch.mean(x, dim=-1, keepdim=True)
+    # s_zm = s - torch.mean(s, dim=-1, keepdim=True)
+    # t = torch.sum(
+        # x_zm * s_zm, dim=-1,
+        # keepdim=True) * s_zm / (l2norm(s_zm, keepdim=True)**2 + eps)
+    # loss = -20. * torch.log10(eps + l2norm(t) / (l2norm(x_zm - t) + eps))
+    # return torch.sum(loss) / x.shape[0]
+
+LAMBDA_WAV = 100
+LAMBDA_ADV = 1
+LAMBDA_REC = 1
+LAMBDA_COM = 1000
+LAMBDA_FEAT = 1
+discriminator_iter_start = 500
+def reconstruction_loss(x, G_x, eps=1e-7):
+    # NOTE (lsx): hard-coded now
+    L = LAMBDA_WAV * F.mse_loss(x, G_x)  # wav L1 loss
+    # loss_sisnr = sisnr_loss(G_x, x) #
+    # L += 0.01*loss_sisnr
+    # 2^6=64 -> 2^10=1024
+    # NOTE (lsx): add 2^11
+    for i in range(6, 12):
+        # for i in range(5, 12): # Encodec setting
+        s = 2**i
+        melspec = MelSpectrogram(
+            sample_rate=16000,
+            n_fft=max(s, 512),
+            win_length=s,
+            hop_length=s // 4,
+            n_mels=64,
+            wkwargs={"device": G_x.device}).to(G_x.device)
+        S_x = melspec(x)
+        S_G_x = melspec(G_x)
+        l1_loss = (S_x - S_G_x).abs().mean()
+        l2_loss = (((torch.log(S_x.abs() + eps) - torch.log(S_G_x.abs() + eps))**2).mean(dim=-2)**0.5).mean()
+
+        alpha = (s / 2) ** 0.5
+        L += (l1_loss + alpha * l2_loss)
+    return L
+
+
+def criterion_d(y_disc_r, y_disc_gen, fmap_r_det, fmap_gen_det, y_df_hat_r,
+                y_df_hat_g, fmap_f_r, fmap_f_g, y_ds_hat_r, y_ds_hat_g,
+                fmap_s_r, fmap_s_g):
+    """Hinge Loss"""
+    loss = 0.0
+    loss1 = 0.0
+    loss2 = 0.0
+    loss3 = 0.0
+    for i in range(len(y_disc_r)):
+        loss1 += F.relu(1 - y_disc_r[i]).mean() + F.relu(1 + y_disc_gen[
+            i]).mean()
+    for i in range(len(y_df_hat_r)):
+        loss2 += F.relu(1 - y_df_hat_r[i]).mean() + F.relu(1 + y_df_hat_g[
+            i]).mean()
+    for i in range(len(y_ds_hat_r)):
+        loss3 += F.relu(1 - y_ds_hat_r[i]).mean() + F.relu(1 + y_ds_hat_g[
+            i]).mean()
+
+    loss = (loss1 / len(y_disc_gen) + loss2 / len(y_df_hat_r) + loss3 /
+            len(y_ds_hat_r)) / 3.0
+
+    return loss
+
+
+def criterion_g(commit_loss, x, G_x, fmap_r, fmap_gen, y_disc_r, y_disc_gen,
+                y_df_hat_r, y_df_hat_g, fmap_f_r, fmap_f_g, y_ds_hat_r,
+                y_ds_hat_g, fmap_s_r, fmap_s_g, args):
+    adv_g_loss = adversarial_g_loss(y_disc_gen)
+    feat_loss = (feature_loss(fmap_r, fmap_gen) + sim_loss(
+        y_disc_r, y_disc_gen) + feature_loss(fmap_f_r, fmap_f_g) + sim_loss(
+            y_df_hat_r, y_df_hat_g) + feature_loss(fmap_s_r, fmap_s_g) +
+                 sim_loss(y_ds_hat_r, y_ds_hat_g)) / 3.0
+    rec_loss = reconstruction_loss(x.contiguous(), G_x.contiguous(), args)
+    total_loss = args.LAMBDA_COM * commit_loss + args.LAMBDA_ADV * adv_g_loss + args.LAMBDA_FEAT * feat_loss + args.LAMBDA_REC * rec_loss
+    return total_loss, adv_g_loss, feat_loss, rec_loss
+
+
+def adopt_weight(weight, global_step, threshold=0, value=0.):
+    if global_step < threshold:
+        weight = value
+    return weight
+
+
+def adopt_dis_weight(weight, global_step, threshold=0, value=0.):
+    # 0,3,6,9,13....这些时间步，不更新dis
+    if global_step % 3 == 0:
+        weight = value
+    return weight
+
+
+def calculate_adaptive_weight(nll_loss, g_loss, last_layer, args):
+    if last_layer is not None:
+        nll_grads = torch.autograd.grad(
+            nll_loss, last_layer, retain_graph=True)[0]
+        g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+    else:
+        print('last_layer cannot be none')
+        assert 1 == 2
+    d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+    d_weight = torch.clamp(d_weight, 1.0, 1.0).detach()
+    d_weight = d_weight * args.LAMBDA_ADV
+    return d_weight
+
+def loss_g(codebook_loss,
+           inputs,
+           reconstructions,
+           fmap_r,
+           fmap_gen,
+           y_disc_r,
+           y_disc_gen,
+           global_step,
+           y_df_hat_r,
+           y_df_hat_g,
+           y_ds_hat_r,
+           y_ds_hat_g,
+           fmap_f_r,
+           fmap_f_g,
+           fmap_s_r,
+           fmap_s_g,
+           last_layer=None,
+           is_training=True,
+           args=None):
+    """
+    args:
+        codebook_loss: commit loss.
+        inputs: ground-truth wav.
+        reconstructions: reconstructed wav.
+        fmap_r: real stft-D feature map.
+        fmap_gen: fake stft-D feature map.
+        y_disc_r: real stft-D logits.
+        y_disc_gen: fake stft-D logits.
+        global_step: global training step.
+        y_df_hat_r: real MPD logits.
+        y_df_hat_g: fake MPD logits.
+        y_ds_hat_r: real MSD logits.
+        y_ds_hat_g: fake MSD logits.
+        fmap_f_r: real MPD feature map.
+        fmap_f_g: fake MPD feature map.
+        fmap_s_r: real MSD feature map.
+        fmap_s_g: fake MSD feature map.
+    """
+    rec_loss = reconstruction_loss(inputs.contiguous(),
+                                   reconstructions.contiguous())
+    adv_g_loss = adversarial_g_loss(y_disc_gen)
+    adv_mpd_loss = adversarial_g_loss(y_df_hat_g)
+    adv_msd_loss = adversarial_g_loss(y_ds_hat_g)
+    adv_loss = (adv_g_loss + adv_mpd_loss + adv_msd_loss
+                ) / 3.0  # NOTE(lsx): need to divide by 3?
+    feat_loss = feature_loss(
+        fmap_r,
+        fmap_gen)  #+ sim_loss(y_disc_r, y_disc_gen) # NOTE(lsx): need logits?
+    feat_loss_mpd = feature_loss(fmap_f_r,
+                                 fmap_f_g)  #+ sim_loss(y_df_hat_r, y_df_hat_g)
+    feat_loss_msd = feature_loss(fmap_s_r,
+                                 fmap_s_g)  #+ sim_loss(y_ds_hat_r, y_ds_hat_g)
+    feat_loss_tot = (feat_loss + feat_loss_mpd + feat_loss_msd) / 3.0
+    d_weight = torch.tensor(1.0)
+    # try:
+    #     d_weight = calculate_adaptive_weight(rec_loss, adv_g_loss, last_layer, args) # 动态调整重构损失和对抗损失
+    # except RuntimeError:
+    #     assert not is_training
+    #     d_weight = torch.tensor(0.0)
+    disc_factor = adopt_weight(
+        LAMBDA_ADV, global_step, threshold=discriminator_iter_start)
+    if disc_factor == 0.:
+        fm_loss_wt = 0
+    else:
+        fm_loss_wt = LAMBDA_FEAT
+    #feat_factor = adopt_weight(args.LAMBDA_FEAT, global_step, threshold=args.discriminator_iter_start)
+    loss = rec_loss + d_weight * disc_factor * adv_loss + \
+           fm_loss_wt * feat_loss_tot + LAMBDA_COM * codebook_loss.mean()
+    return loss, rec_loss, adv_loss, feat_loss_tot, d_weight
+
+
+def loss_dis(y_disc_r_det, y_disc_gen_det, fmap_r_det, fmap_gen_det, y_df_hat_r,
+             y_df_hat_g, fmap_f_r, fmap_f_g, y_ds_hat_r, y_ds_hat_g, fmap_s_r,
+             fmap_s_g, global_step):
+    disc_factor = adopt_weight(
+        LAMBDA_ADV, global_step, threshold=discriminator_iter_start)
+    d_loss = disc_factor * criterion_d(y_disc_r_det, y_disc_gen_det, fmap_r_det,
+                                       fmap_gen_det, y_df_hat_r, y_df_hat_g,
+                                       fmap_f_r, fmap_f_g, y_ds_hat_r,
+                                       y_ds_hat_g, fmap_s_r, fmap_s_g)
+    return d_loss
+
+class AttentionCTCLoss(torch.nn.Module):
+    def __init__(self, blank_logprob=-1):
+        super(AttentionCTCLoss, self).__init__()
+        self.log_softmax = torch.nn.LogSoftmax(dim=3)
+        self.blank_logprob = blank_logprob
+        self.CTCLoss = torch.nn.CTCLoss(zero_infinity=True)
+
+    def forward(self, attn_logprob, in_lens, out_lens):
+        key_lens = in_lens
+        query_lens = out_lens
+        attn_logprob_padded = F.pad(
+            input=attn_logprob, pad=(1, 0, 0, 0, 0, 0, 0, 0),
+            value=self.blank_logprob)
+        cost_total = 0.0
+        for bid in range(attn_logprob.shape[0]):
+            target_seq = torch.arange(1, key_lens[bid]+1).unsqueeze(0)
+            curr_logprob = attn_logprob_padded[bid].permute(1, 0, 2)[
+                :query_lens[bid], :, :key_lens[bid]+1]
+            curr_logprob = self.log_softmax(curr_logprob[None])[0]
+            ctc_cost = self.CTCLoss(curr_logprob, target_seq,
+                                    input_lengths=query_lens[bid:bid+1],
+                                    target_lengths=key_lens[bid:bid+1])
+            cost_total += ctc_cost
+        cost = cost_total/attn_logprob.shape[0]
+        return cost
+
+
+class FocalLoss(torch.nn.Module):
+
+    def __init__(self, gamma=0, eps=1e-7):
+        super(FocalLoss, self).__init__()
+        self.gamma = gamma
+        self.eps = eps
+        self.ce = torch.nn.CrossEntropyLoss()
+
+    def forward(self, input, target):
+        logp = self.ce(input, target)
+        p = torch.exp(-logp)
+        loss = (1 - p) ** self.gamma * logp
+        return loss.mean()
+
+def feature_loss(fmap_r, fmap_g):
+    loss = 0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss += torch.mean(torch.abs(rl - gl))
+
+    return loss * 2
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    r_losses = []
+    g_losses = []
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        r_loss = torch.mean((1 - dr) ** 2)
+        g_loss = torch.mean(dg ** 2)
+        loss += (r_loss + g_loss)
+        r_losses.append(r_loss.item())
+        g_losses.append(g_loss.item())
+
+    return loss, r_losses, g_losses
+
+
+def generator_loss(disc_outputs):
+    loss = 0
+    gen_losses = []
+    for dg in disc_outputs:
+        l = torch.mean((1 - dg) ** 2)
+        gen_losses.append(l)
+        loss += l
+
+    return loss, gen_losses

meldataset.py

@@ -0,0 +1,131 @@
+# coding: utf-8
+import os
+import os.path as osp
+import time
+import random
+import numpy as np
+import random
+import soundfile as sf
+import librosa
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torchaudio
+from torch.utils.data import DataLoader
+
+import math
+
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+from torch.utils.data.distributed import DistributedSampler
+
+
+np.random.seed(114514)
+random.seed(114514)
+SPECT_PARAMS = {
+    "n_fft": 2048,
+    "win_length": 1200,
+    "hop_length": 300,
+}
+MEL_PARAMS = {
+    "n_mels": 80,
+}
+
+to_mel = torchaudio.transforms.MelSpectrogram(
+    n_mels=MEL_PARAMS['n_mels'], **SPECT_PARAMS)
+mean, std = -4, 4
+
+
+def preprocess(wave):
+    # wave = wave.unsqueeze(0)
+    wave_tensor = torch.from_numpy(wave).float()
+    mel_tensor = to_mel(wave_tensor)
+    mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std
+    return mel_tensor
+
+
+class PseudoDataset(torch.utils.data.Dataset):
+    def __init__(self,
+                 list_path,
+                 sr=24000,
+                 range=(1, 30), # length of the audio duration in seconds
+                 ):
+
+        self.data_list = [] # read your list path here
+        self.sr = sr
+        self.duration_range = range
+
+    def __len__(self):
+        # return len(self.data_list)
+        return 100 # return a fixed number for testing
+
+    def __getitem__(self, idx):
+        # replace this with your own data loading
+        # wave, sr = librosa.load(self.data_list[idx], sr=self.sr)
+        wave = np.random.randn(self.sr * random.randint(*self.duration_range)).clamp(-1, 1)
+        mel = preprocess(wave)
+        return wave, mel
+
+
+def collate(batch):
+    # batch[0] = wave, mel, text, f0, speakerid
+    batch_size = len(batch)
+
+    # sort by mel length
+    lengths = [b[1].shape[1] for b in batch]
+    batch_indexes = np.argsort(lengths)[::-1]
+    batch = [batch[bid] for bid in batch_indexes]
+
+    nmels = batch[0][1].size(0)
+    max_mel_length = max([b[1].shape[1] for b in batch])
+    max_wave_length = max([b[0].size(0) for b in batch])
+
+    mels = torch.zeros((batch_size, nmels, max_mel_length)).float() - 10
+    waves = torch.zeros((batch_size, max_wave_length)).float()
+
+    mel_lengths = torch.zeros(batch_size).long()
+    wave_lengths = torch.zeros(batch_size).long()
+
+    for bid, (wave, mel) in enumerate(batch):
+        mel_size = mel.size(1)
+        mels[bid, :, :mel_size] = mel
+        waves[bid, : wave.size(0)] = wave
+        mel_lengths[bid] = mel_size
+        wave_lengths[bid] = wave.size(0)
+
+    return waves, mels, wave_lengths, mel_lengths
+
+
+def build_dataloader(
+    rank=0,
+    world_size=1,
+    batch_size=32,
+    num_workers=0,
+    prefetch_factor=16,
+):
+    dataset = PseudoDataset() # replace this with your own dataset
+    collate_fn = collate
+    sampler = torch.utils.data.distributed.DistributedSampler(
+        dataset,
+        num_replicas=world_size,
+        rank=rank,
+        shuffle=True,
+        seed=114514,
+    )
+    data_loader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        sampler=sampler,
+        num_workers=num_workers,
+        drop_last=True,
+        collate_fn=collate_fn,
+        pin_memory=True,
+        prefetch_factor=prefetch_factor,
+        # shuffle=True,
+    )
+
+    return data_loader
+

optimizers.py

@@ -0,0 +1,108 @@
+#coding:utf-8
+import os, sys
+import os.path as osp
+import numpy as np
+import torch
+from torch import nn
+from torch.optim import Optimizer
+from functools import reduce
+from torch.optim import AdamW
+
+class MultiOptimizer:
+    def __init__(self, optimizers={}, schedulers={}):
+        self.optimizers = optimizers
+        self.schedulers = schedulers
+        self.keys = list(optimizers.keys())
+        self.param_groups = reduce(lambda x,y: x+y, [v.param_groups for v in self.optimizers.values()])
+
+    def state_dict(self):
+        state_dicts = [(key, self.optimizers[key].state_dict())\
+                       for key in self.keys]
+        return state_dicts
+
+    def scheduler_state_dict(self):
+        state_dicts = [(key, self.schedulers[key].state_dict())\
+                       for key in self.keys]
+        return state_dicts
+
+    def load_state_dict(self, state_dict):
+        for key, val in state_dict:
+            try:
+                self.optimizers[key].load_state_dict(val)
+            except:
+                print("Unloaded %s" % key)
+
+    def load_scheduler_state_dict(self, state_dict):
+        for key, val in state_dict:
+            try:
+                self.schedulers[key].load_state_dict(val)
+            except:
+                print("Unloaded %s" % key)
+
+    def step(self, key=None, scaler=None):
+        keys = [key] if key is not None else self.keys
+        _ = [self._step(key, scaler) for key in keys]
+
+    def _step(self, key, scaler=None):
+        if scaler is not None:
+            scaler.step(self.optimizers[key])
+            scaler.update()
+        else:
+            self.optimizers[key].step()
+
+    def zero_grad(self, key=None):
+        if key is not None:
+            self.optimizers[key].zero_grad()
+        else:
+            _ = [self.optimizers[key].zero_grad() for key in self.keys]
+
+    def scheduler(self, *args, key=None):
+        if key is not None:
+            self.schedulers[key].step(*args)
+        else:
+            _ = [self.schedulers[key].step_batch(*args) for key in self.keys]
+
+def define_scheduler(optimizer, params):
+    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=params['gamma'])
+
+    return scheduler
+
+from transformer_modules.optim import Eden, ScaledAdam
+
+def build_optimizer(model_dict, scheduler_params_dict, lr, type='AdamW'):
+    optim = {}
+    for key, model in model_dict.items():
+        model_parameters = model.parameters()
+        parameters_names = []
+        parameters_names.append(
+            [
+                name_param_pair[0]
+                for name_param_pair in model.named_parameters()
+            ]
+        )
+        if type == 'ScaledAdam':
+            optim[key] = ScaledAdam(
+                model_parameters,
+                lr=lr,
+                betas=(0.9, 0.95),
+                clipping_scale=2.0,
+                parameters_names=parameters_names,
+                show_dominant_parameters=False,
+                clipping_update_period=1000,
+            )
+        elif type == 'AdamW':
+            optim[key] = AdamW(
+                model_parameters,
+                lr=lr,
+                betas=(0.9, 0.98),
+                eps=1e-9,
+                weight_decay=0.1,
+            )
+        else:
+            raise ValueError('Unknown optimizer type: %s' % type)
+
+    schedulers = dict([(key, torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.999996))
+                       for key, opt in optim.items()])
+
+    multi_optim = MultiOptimizer(optim, schedulers)
+    return multi_optim