From 0a49d1ee002a7874300f89eb043516ec3da067e5 Mon Sep 17 00:00:00 2001 From: litagin02 Date: Fri, 9 Feb 2024 17:31:27 +0900 Subject: [PATCH] Fix: wavlm shape bug when batch_size 1 --- losses.py | 4 +++- train_ms_jp_extra.py | 8 +++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/losses.py b/losses.py index 62982fc..763cc02 100644 --- a/losses.py +++ b/losses.py @@ -2,6 +2,8 @@ import torch import torchaudio from transformers import AutoModel +from common.log import logger + def feature_loss(fmap_r, fmap_g): loss = 0 @@ -78,7 +80,7 @@ class WavLMLoss(torch.nn.Module): ).hidden_states y_rec_16 = self.resample(y_rec) y_rec_embeddings = self.wavlm( - input_values=y_rec_16.squeeze(), output_hidden_states=True + input_values=y_rec_16, output_hidden_states=True ).hidden_states floss = 0 diff --git a/train_ms_jp_extra.py b/train_ms_jp_extra.py index e20868f..2c46eb2 100644 --- a/train_ms_jp_extra.py +++ b/train_ms_jp_extra.py @@ -717,9 +717,11 @@ def train_and_evaluate( ) scaler.step(optim_dur_disc) if net_wd is not None: + # logger.debug(f"y.shape: {y.shape}, y_hat.shape: {y_hat.shape}") + # shape: (batch, 1, time) with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16): loss_slm = wl.discriminator( - y.detach().squeeze(), y_hat.detach().squeeze() + y.detach().squeeze(1), y_hat.detach().squeeze(1) ).mean() optim_wd.zero_grad() @@ -743,8 +745,8 @@ def train_and_evaluate( if net_dur_disc is not None: _, y_dur_hat_g = net_dur_disc(hidden_x, x_mask, logw_, logw, g) if net_wd is not None: - loss_lm = wl(y.detach().squeeze(), y_hat.squeeze()).mean() - loss_lm_gen = wl.generator(y_hat.squeeze()) + loss_lm = wl(y.detach().squeeze(1), y_hat.squeeze(1)).mean() + loss_lm_gen = wl.generator(y_hat.squeeze(1)) with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16): loss_dur = torch.sum(l_length.float()) loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel