tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Fix: wavlm shape bug when batch_size 1

Browse Source
This commit is contained in:
litagin02
2024-02-09 17:31:27 +09:00
parent fd1d3efd18
commit 0a49d1ee00
2 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
losses.py
View File

@@ -2,6 +2,8 @@ import torch
import torchaudio import torchaudio
from transformers import AutoModel from transformers import AutoModel
from common.log import logger
def feature_loss(fmap_r, fmap_g): def feature_loss(fmap_r, fmap_g):
loss = 0 loss = 0
@@ -78,7 +80,7 @@ class WavLMLoss(torch.nn.Module):
).hidden_states ).hidden_states
y_rec_16 = self.resample(y_rec) y_rec_16 = self.resample(y_rec)
y_rec_embeddings = self.wavlm( 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 ).hidden_states
floss = 0 floss = 0

8
train_ms_jp_extra.py
View File

@@ -717,9 +717,11 @@ def train_and_evaluate(
) )
scaler.step(optim_dur_disc) scaler.step(optim_dur_disc)
if net_wd is not None: 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): with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
loss_slm = wl.discriminator( loss_slm = wl.discriminator(
y.detach().squeeze(), y_hat.detach().squeeze() y.detach().squeeze(1), y_hat.detach().squeeze(1)
).mean() ).mean()
optim_wd.zero_grad() optim_wd.zero_grad()
@@ -743,8 +745,8 @@ def train_and_evaluate(
if net_dur_disc is not None: if net_dur_disc is not None:
_, y_dur_hat_g = net_dur_disc(hidden_x, x_mask, logw_, logw, g) _, y_dur_hat_g = net_dur_disc(hidden_x, x_mask, logw_, logw, g)
if net_wd is not None: if net_wd is not None:
loss_lm = wl(y.detach().squeeze(), y_hat.squeeze()).mean() loss_lm = wl(y.detach().squeeze(1), y_hat.squeeze(1)).mean()
loss_lm_gen = wl.generator(y_hat.squeeze()) loss_lm_gen = wl.generator(y_hat.squeeze(1))
with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16): with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
loss_dur = torch.sum(l_length.float()) loss_dur = torch.sum(l_length.float())
loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel