VAE Training Needed

[yukiarimo]

Hello!

Can someone please implement training loop & loss function for VAE in full fine-tuning (both encoder-decoder & decoder-only scenarios)?

Thanks!

[a710128]

Our training code is based on DAC. And see #145 (comment) for detailed information on hyperparameters.

[yukiarimo]

I've tried train it like this:

#!/usr/bin/env python3
import json
import argparse
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import soundfile as sf
import torchaudio
from pathlib import Path
from torch.utils.data import DataLoader, Dataset
from torch.optim import AdamW
from torch.nn.utils import weight_norm, spectral_norm

from voxcpm.modules.audiovae import AudioVAE, AudioVAEConfig
from voxcpm.training import Accelerator

# ==========================================
# 1. DAC DISCRIMINATOR ARCHITECTURE (Ported)
# Source: discriminator.py (Simplified)
# ==========================================


def WNConv1d(*args, **kwargs):
    return weight_norm(nn.Conv1d(*args, **kwargs))


def WNConv2d(*args, **kwargs):
    return weight_norm(nn.Conv2d(*args, **kwargs))


class DiskDiscriminatorP(nn.Module):
    """Multi-Period Discriminator (MPD)"""

    def __init__(self, period, kernel_size=5, stride=3):
        super().__init__()
        self.period = period
        self.convs = nn.ModuleList(
            [
                WNConv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(2, 0)),
                WNConv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(2, 0)),
                WNConv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(2, 0)),
                WNConv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(2, 0)),
                WNConv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0)),
            ]
        )
        self.conv_post = WNConv2d(1024, 1, (3, 1), 1, padding=(1, 0))

    def forward(self, x):
        fmap = []
        b, c, t = x.shape
        if t % self.period != 0:
            n_pad = self.period - (t % self.period)
            x = F.pad(x, (0, n_pad), "reflect")
            t = x.shape[-1]
        x = x.view(b, c, t // self.period, self.period)

        for l in self.convs:
            x = l(x)
            x = F.leaky_relu(x, 0.1)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        return torch.flatten(x, 1, -1), fmap


class DiskDiscriminatorS(nn.Module):
    """Multi-Scale Discriminator (MSD)"""

    def __init__(self):
        super().__init__()
        self.convs = nn.ModuleList(
            [
                WNConv1d(1, 16, 15, 1, padding=7),
                WNConv1d(16, 64, 41, 4, groups=4, padding=20),
                WNConv1d(64, 256, 41, 4, groups=16, padding=20),
                WNConv1d(256, 1024, 41, 4, groups=64, padding=20),
                WNConv1d(1024, 1024, 41, 4, groups=256, padding=20),
                WNConv1d(1024, 1024, 5, 1, padding=2),
            ]
        )
        self.conv_post = WNConv1d(1024, 1, 3, 1, padding=1)

    def forward(self, x):
        fmap = []
        for l in self.convs:
            x = l(x)
            x = F.leaky_relu(x, 0.1)
            fmap.append(x)
        x = self.conv_post(x)
        fmap.append(x)
        return torch.flatten(x, 1, -1), fmap


class DACDiscriminator(nn.Module):
    def __init__(self):
        super().__init__()
        # MPD Periods from DAC config
        self.mpds = nn.ModuleList([DiskDiscriminatorP(p) for p in [2, 3, 5, 7, 11]])
        # MSDs (Original + Downsampled)
        self.msds = nn.ModuleList(
            [DiskDiscriminatorS(), DiskDiscriminatorS(), DiskDiscriminatorS()]
        )
        self.downsample = torchaudio.transforms.Resample(
            48000, 24000
        )  # Simple pooling approx via resample

    def forward(self, x):
        # MPD
        y_ds, fmaps = [], []
        for d in self.mpds:
            y, f = d(x)
            y_ds.append(y)
            fmaps.append(f)

        # MSD (Scale 1)
        y, f = self.msds[0](x)
        y_ds.append(y)
        fmaps.append(f)

        # MSD (Scale 2 - Downsampled)
        x_down = F.avg_pool1d(x, 2)
        y, f = self.msds[1](x_down)
        y_ds.append(y)
        fmaps.append(f)

        # MSD (Scale 3 - Downsampled again)
        x_down2 = F.avg_pool1d(x_down, 2)
        y, f = self.msds[2](x_down2)
        y_ds.append(y)
        fmaps.append(f)

        return y_ds, fmaps


# ==========================================
# 2. LOSS FUNCTIONS (Exact Specs)
# ==========================================


class DACLosses(nn.Module):
    def __init__(self, sample_rate=48000):
        super().__init__()
        # Official hyperparameters provided by developer
        self.n_mels = [5, 10, 20, 40, 80, 160, 320]
        self.window_lengths = [32, 64, 128, 256, 512, 1024, 2048]
        self.mel_transforms = nn.ModuleList()

        for w, n_mel in zip(self.window_lengths, self.n_mels):
            hop = w // 4
            self.mel_transforms.append(
                torchaudio.transforms.MelSpectrogram(
                    sample_rate=sample_rate,
                    n_fft=w,
                    win_length=w,
                    hop_length=hop,
                    n_mels=n_mel,
                    power=1.0,
                    normalized=True,
                    center=True,  # 'slaney' norm approx
                )
            )

    def mel_loss(self, x_rec, x_real):
        loss = 0.0
        for i, trans in enumerate(self.mel_transforms):
            # Move transform to device dynamically
            if trans.spectrogram.window.device != x_rec.device:
                trans = trans.to(x_rec.device)

            real_mel = trans(x_real.squeeze(1))
            rec_mel = trans(x_rec.squeeze(1))

            # Log Mel Loss (L1 on Log)
            # clamp_eps: 1.0e-5
            loss += F.l1_loss(
                torch.log(torch.clamp(real_mel, min=1e-5)),
                torch.log(torch.clamp(rec_mel, min=1e-5)),
            )
        return loss

    def feature_loss(self, 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

    def generator_loss(self, disc_outputs):
        loss = 0
        for dg in disc_outputs:
            loss += (
                torch.mean(torch.relu(1 - dg)) ** 2
            )  # LSGAN style often used, or Hinge
        return loss

    def discriminator_loss(self, disc_real, disc_fake):
        loss = 0
        for dr, df in zip(disc_real, disc_fake):
            loss += torch.mean(torch.relu(1 - dr)) ** 2
            loss += torch.mean(torch.relu(1 + df)) ** 2
        return loss


# ==========================================
# 3. ROBUST DATA LOADING
# ==========================================


def robust_load(path, target_sr=48000):
    try:
        wav_np, sr = sf.read(path)
        if len(wav_np.shape) == 1:
            wav = torch.from_numpy(wav_np).float().unsqueeze(0)
        else:
            wav = torch.from_numpy(wav_np).float().t()
        if wav.shape[0] > 1:
            wav = wav.mean(dim=0, keepdim=True)
        if sr != target_sr:
            wav = torchaudio.functional.resample(wav, sr, target_sr)
        if torch.isnan(wav).any():
            return None
        return wav
    except:
        return None


class VAEAudioDataset(Dataset):
    def __init__(self, manifest_path, sample_rate=48000, max_duration=2.0):
        self.audio_paths = []
        self.sample_rate = sample_rate
        self.max_samples = int(sample_rate * max_duration)
        with open(manifest_path, "r") as f:
            for line in f:
                try:
                    if "{" in line:
                        data = json.loads(line.strip())
                        if "audio" in data:
                            self.audio_paths.append(data["audio"])
                    elif "|" in line:
                        parts = line.split("|")
                        if len(parts) > 0:
                            self.audio_paths.append(parts[0])
                except:
                    pass
        print(f"Loaded {len(self.audio_paths)} files.")

    def __len__(self):
        return len(self.audio_paths)

    def __getitem__(self, idx):
        wav = robust_load(self.audio_paths[idx], self.sample_rate)
        if wav is None:
            return torch.zeros(1, self.max_samples)
        if wav.shape[-1] >= self.max_samples:
            start = np.random.randint(0, wav.shape[-1] - self.max_samples)
            wav = wav[:, start : start + self.max_samples]
        else:
            wav = F.pad(wav, (0, self.max_samples - wav.shape[-1]))
        return wav


def collate_fn(batch):
    filtered = [b for b in batch if b.abs().sum() > 0]
    return torch.stack(filtered) if filtered else torch.tensor([])


# ==========================================
# 4. MAIN TRAINING LOOP
# ==========================================


def train():
    parser = argparse.ArgumentParser()
    parser.add_argument("--train_manifest", type=str, required=True)
    parser.add_argument("--pretrained_path", type=str, required=True)
    parser.add_argument("--save_path", type=str, default="checkpoints_dac_vae")
    parser.add_argument(
        "--mode",
        type=str,
        default="decoder_only",
        choices=["decoder_only", "encoder_decoder"],
    )
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--lr", type=float, default=3e-5)  # Dev provided: 3e-5
    parser.add_argument("--steps", type=int, default=20000)
    parser.add_argument("--save_interval", type=int, default=1000)
    args = parser.parse_args()

    # Setup
    accelerator = Accelerator(amp=False)  # GANs prefer stability over AMP usually
    Path(args.save_path).mkdir(parents=True, exist_ok=True)
    pretrained_dir = Path(args.pretrained_path)

    # Load VAE
    with open(pretrained_dir / "config.json", "r") as f:
        full_config = json.load(f)
    vae = AudioVAE(config=AudioVAEConfig(**full_config["audio_vae_config"]))

    # Load VAE Weights
    ckpt = torch.load(pretrained_dir / "audiovae.pth", map_location="cpu")
    sd = ckpt["state_dict"] if "state_dict" in ckpt else ckpt
    vae.load_state_dict(
        {k.replace("audio_vae.", ""): v for k, v in sd.items()}, strict=True
    )
    vae = vae.to(accelerator.device)

    # Initialize Discriminator (From Scratch)
    discriminator = DACDiscriminator().to(accelerator.device)

    # Optimizers
    for p in vae.parameters():
        p.requires_grad = False
    trainable_params = []

    if args.mode == "encoder_decoder":
        for p in vae.encoder.parameters():
            p.requires_grad = True
            trainable_params.append(p)
        # Assuming VAE specific params (proj) are in encoder for simplicity, or we add them specifically
    for p in vae.decoder.parameters():
        p.requires_grad = True
        trainable_params.append(p)

    # Dev Param: LR 3e-5
    opt_g = AdamW(trainable_params, lr=args.lr, betas=(0.8, 0.99))
    opt_d = AdamW(discriminator.parameters(), lr=args.lr, betas=(0.8, 0.99))

    # Loss Wrapper
    dac_losses = DACLosses(sample_rate=vae.sample_rate).to(accelerator.device)

    # Data
    ds = VAEAudioDataset(args.train_manifest, sample_rate=vae.sample_rate)
    dl = DataLoader(
        ds,
        batch_size=args.batch_size,
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=4,
    )

    W_MEL = 60.0
    W_FEAT = 2.0
    W_ADV = 0.5
    W_KL = 1e-4

    print("Starting DAC-Style VAE Training...")
    step = 0
    vae.train()
    discriminator.train()

    while step < args.steps:
        for x_real in dl:
            if x_real.numel() == 0:
                continue
            x_real = x_real.to(accelerator.device)

            # --- Forward VAE ---
            if args.mode == "decoder_only":
                with torch.no_grad():
                    # get 'mu' and 'logvar' from encoder internally?
                    # Standard AudioVAE.encode only returns mu.
                    # We need to access internal result for KL loss.
                    # Let's perform manual encoding pass to get distribution
                    enc_out = vae.encoder(vae.preprocess(x_real, vae.sample_rate))
                    mu, logvar = enc_out["mu"], enc_out["logvar"]

                    # Clip Logvar as per dev recommendation [-20, 20]
                    logvar = torch.clamp(logvar, -20, 20)
            else:
                enc_out = vae.encoder(vae.preprocess(x_real, vae.sample_rate))
                mu, logvar = enc_out["mu"], enc_out["logvar"]
                logvar = torch.clamp(logvar, -20, 20)

            # Reparameterize
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std

            x_fake = vae.decode(z)

            # Match lengths
            min_l = min(x_real.shape[-1], x_fake.shape[-1])
            x_real_crop = x_real[..., :min_l]
            x_fake_crop = x_fake[..., :min_l]

            # --- Train Discriminator ---
            opt_d.zero_grad()
            y_real, _ = discriminator(x_real_crop)
            y_fake, _ = discriminator(x_fake_crop.detach())

            loss_d = dac_losses.discriminator_loss(y_real, y_fake)
            loss_d.backward()

            # Dev param: Clip grad 10.0
            torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 10.0)
            opt_d.step()

            # --- Train Generator ---
            opt_g.zero_grad()
            y_fake_g, fmap_fake = discriminator(x_fake_crop)
            _, fmap_real = discriminator(x_real_crop)

            # 1. Mel Loss
            loss_mel = dac_losses.mel_loss(x_fake_crop, x_real_crop)

            # 2. Adv Loss
            loss_adv = dac_losses.generator_loss(y_fake_g)

            # 3. Feature Match
            loss_feat = dac_losses.feature_loss(fmap_real, fmap_fake)

            # 4. KL Loss
            # sum over dimensions, mean over batch
            loss_kl = (
                -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1).mean()
            )

            loss_g = (
                (W_MEL * loss_mel)
                + (W_FEAT * loss_feat)
                + (W_ADV * loss_adv)
                + (W_KL * loss_kl)
            )

            loss_g.backward()

            # Dev param: Clip grad 1000.0
            torch.nn.utils.clip_grad_norm_(trainable_params, 1000.0)
            opt_g.step()

            step += 1
            if step % 20 == 0:
                print(
                    f"Step {step}: D={loss_d.item():.3f} | G={loss_g.item():.3f} [Mel:{loss_mel.item():.3f} Adv:{loss_adv.item():.3f} KL:{loss_kl.item():.5f}]"
                )

            if step % args.save_interval == 0:
                out = Path(args.save_path) / f"step_{step}"
                out.mkdir(parents=True, exist_ok=True)
                torch.save({"state_dict": vae.state_dict()}, out / "audiovae.pth")
                torch.save(discriminator.state_dict(), out / "discriminator.pth")
                print(f"Saved: {out}")

            if step >= args.steps:
                break


if __name__ == "__main__":
    train()

But it doesn't make much difference and just introduces metallic artifacts (audio dataset is studio quality).

Can you give me any suggestions how I can improve this?

[ziyangwang1992]

Could you provide the specific training parameters, especially how many GPUs were used and how many training steps were run? One possible reason is that the discriminator was trained from scratch and hasn't been fully trained yet, which may cause robotic/electronic artifacts. Another possible reason is that the hyperparameters haven't been tuned properly, such as the KL weight, learning rate, etc., which may need further adjustment.