lt-tensor 0.0.1a33__py3-none-any.whl → 0.0.1a35__py3-none-any.whl

lt_tensor/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.0.1a33"
+__version__ = "0.0.1a35"
 
 from . import (
     lr_schedulers,

lt_tensor/losses.py

@@ -6,19 +6,24 @@ __all__ = [
     "hybrid_loss",
     "diff_loss",
     "cosine_loss",
-    "gan_loss",
     "ft_n_loss",
+    "MultiMelScaleLoss",
 ]
 import math
 import random
 from lt_tensor.torch_commons import *
 from lt_utils.common import *
 import torch.nn.functional as F
+from lt_tensor.model_base import Model
+from lt_tensor.processors import AudioProcessor, AudioProcessorConfig
+from lt_tensor.math_ops import normalize_minmax, normalize_unit_norm, normalize_zscore
+
 
 def ft_n_loss(output: Tensor, target: Tensor, weight: Optional[Tensor] = None):
     if weight is not None:
-        return torch.mean((torch.abs(output - target) + weight) **0.5)
-    return torch.mean(torch.abs(output - target)**0.5)
+        return torch.mean((torch.abs(output - target) + weight) ** 0.5)
+    return torch.mean(torch.abs(output - target) ** 0.5)
+
 
 def adaptive_l1_loss(
     inp: Tensor,
@@ -58,50 +63,6 @@ def cosine_loss(inp, tgt):
     return 1 - cos.mean()  # Lower is better
 
 
-class GanLosses:
-    @staticmethod
-    def get_loss(
-        pred: Tensor,
-        target_is_real: bool,
-        loss_type: Literal["bce", "mse", "hinge", "wasserstein"] = "bce",
-    ) -> Tensor:
-        if loss_type == "bce":  # Standard GAN
-            target = torch.ones_like(pred) if target_is_real else torch.zeros_like(pred)
-            return F.binary_cross_entropy_with_logits(pred, target)
-
-        elif loss_type == "mse":  # LSGAN
-            target = torch.ones_like(pred) if target_is_real else torch.zeros_like(pred)
-            return F.mse_loss(torch.sigmoid(pred), target)
-
-        elif loss_type == "hinge":
-            if target_is_real:
-                return torch.mean(F.relu(1.0 - pred))
-            else:
-                return torch.mean(F.relu(1.0 + pred))
-
-        elif loss_type == "wasserstein":
-            return -pred.mean() if target_is_real else pred.mean()
-
-        else:
-            raise ValueError(f"Unknown loss_type: {loss_type}")
-
-    @staticmethod
-    def generator_loss(fake_pred: Tensor, loss_type: str = "bce") -> Tensor:
-        return GanLosses.get_loss(fake_pred, target_is_real=True, loss_type=loss_type)
-
-    @staticmethod
-    def discriminator_loss(
-        real_pred: Tensor, fake_pred: Tensor, loss_type: str = "bce"
-    ) -> Tensor:
-        real_loss = GanLosses.get_loss(
-            real_pred, target_is_real=True, loss_type=loss_type
-        )
-        fake_loss = GanLosses.get_loss(
-            fake_pred.detach(), target_is_real=False, loss_type=loss_type
-        )
-        return (real_loss + fake_loss) * 0.5
-
-
 def masked_cross_entropy(
     logits: torch.Tensor,  # [B, T, V]
     targets: torch.Tensor,  # [B, T]
@@ -157,3 +118,164 @@ def gan_d_loss(real_preds, fake_preds, use_lsgan=True):
                 torch.log(1 - fake + 1e-7)
             )
     return loss
+
+
+class MultiMelScaleLoss(Model):
+    def __init__(
+        self,
+        sample_rate: int,
+        n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320],
+        window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
+        n_ffts: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
+        hops: List[int] = [8, 16, 32, 64, 128, 256, 512],
+        f_min: float = [0, 0, 0, 0, 0, 0, 0],
+        f_max: Optional[float] = [None, None, None, None, None, None, None],
+        loss_mel_fn: Callable[[Tensor, Tensor], Tensor] = nn.L1Loss(),
+        loss_pitch_fn: Callable[[Tensor, Tensor], Tensor] = nn.L1Loss(),
+        loss_rms_fn: Callable[[Tensor, Tensor], Tensor] = nn.L1Loss(),
+        center: bool = True,
+        power: float = 1.0,
+        normalized: bool = False,
+        pad_mode: str = "reflect",
+        onesided: Optional[bool] = None,
+        std: int = 4,
+        mean: int = -4,
+        use_istft_norm: bool = True,
+        use_pitch_loss: bool = True,
+        use_rms_loss: bool = True,
+        norm_pitch_fn: Callable[[Tensor], Tensor] = normalize_minmax,
+        norm_rms_fn: Callable[[Tensor], Tensor] = normalize_zscore,
+        lambda_mel: float = 1.0,
+        lambda_rms: float = 1.0,
+        lambda_pitch: float = 1.0,
+        weight: float = 1.0,
+    ):
+        super().__init__()
+        assert (
+            len(n_mels)
+            == len(window_lengths)
+            == len(n_ffts)
+            == len(hops)
+            == len(f_min)
+            == len(f_max)
+        )
+        self.loss_mel_fn = loss_mel_fn
+        self.loss_pitch_fn = loss_pitch_fn
+        self.loss_rms_fn = loss_rms_fn
+        self.lambda_mel = lambda_mel
+        self.weight = weight
+        self.use_istft_norm = use_istft_norm
+        self.use_pitch_loss = use_pitch_loss
+        self.use_rms_loss = use_rms_loss
+        self.lambda_pitch = lambda_pitch
+        self.lambda_rms = lambda_rms
+
+        self.norm_pitch_fn = norm_pitch_fn
+        self.norm_rms = norm_rms_fn
+
+        self._setup_mels(
+            sample_rate,
+            n_mels,
+            window_lengths,
+            n_ffts,
+            hops,
+            f_min,
+            f_max,
+            center,
+            power,
+            normalized,
+            pad_mode,
+            onesided,
+            std,
+            mean,
+        )
+
+    def _setup_mels(
+        self,
+        sample_rate: int,
+        n_mels: List[int],
+        window_lengths: List[int],
+        n_ffts: List[int],
+        hops: List[int],
+        f_min: List[float],
+        f_max: List[Optional[float]],
+        center: bool,
+        power: float,
+        normalized: bool,
+        pad_mode: str,
+        onesided: Optional[bool],
+        std: int,
+        mean: int,
+    ):
+        assert (
+            len(n_mels)
+            == len(window_lengths)
+            == len(n_ffts)
+            == len(hops)
+            == len(f_min)
+            == len(f_max)
+        )
+        _mel_kwargs = dict(
+            sample_rate=sample_rate,
+            center=center,
+            onesided=onesided,
+            normalized=normalized,
+            power=power,
+            pad_mode=pad_mode,
+            std=std,
+            mean=mean,
+        )
+        self.mel_spectrograms: List[AudioProcessor] = nn.ModuleList(
+            [
+                AudioProcessor(
+                    AudioProcessorConfig(
+                        **_mel_kwargs,
+                        n_mels=mel,
+                        n_fft=n_fft,
+                        win_length=win,
+                        hop_length=hop,
+                        f_min=fmin,
+                        f_max=fmax,
+                    )
+                )
+                for mel, win, n_fft, hop, fmin, fmax in zip(
+                    n_mels, window_lengths, n_ffts, hops, f_min, f_max
+                )
+            ]
+        )
+
+    def forward(
+        self, input_wave: torch.Tensor, target_wave: torch.Tensor
+    ) -> torch.Tensor:
+        assert self.use_istft_norm or input_wave.shape[-1] == target_wave.shape[-1]
+        target_wave = target_wave.to(input_wave.device)
+        losses = 0.0
+        for M in self.mel_spectrograms:
+            # Apply normalization if requested
+            if self.use_istft_norm:
+                input_proc = M.istft_norm(input_wave, length=target_wave.shape[-1])
+                target_proc = M.istft_norm(target_wave, length=target_wave.shape[-1])
+            else:
+                input_proc, target_proc = input_wave, target_wave
+
+            x_mels = M(input_proc)
+            y_mels = M(target_proc)
+
+            loss = self.loss_mel_fn(x_mels.squeeze(), y_mels.squeeze())
+            losses += loss * self.lambda_mel
+
+            # pitch/f0 loss
+            if self.use_pitch_loss:
+                x_pitch = self.norm_pitch_fn(M.compute_pitch(input_proc))
+                y_pitch = self.norm_pitch_fn(M.compute_pitch(target_proc))
+                f0_loss = self.loss_pitch_fn(x_pitch, y_pitch)
+                losses += f0_loss * self.lambda_pitch
+
+            # energy/rms loss
+            if self.use_rms_loss:
+                x_rms = self.norm_rms(M.compute_rms(input_proc, x_mels))
+                y_rms = self.norm_rms(M.compute_rms(target_proc, y_mels))
+                rms_loss = self.loss_rms_fn(x_rms, y_rms)
+                losses += rms_loss * self.lambda_rms
+
+        return losses * self.weight

lt_tensor/lr_schedulers.py

@@ -1,15 +1,20 @@
 __all__ = [
     "WarmupDecayScheduler",
     "AdaptiveDropScheduler",
-    "WaveringLRScheduler",
+    "SinusoidalDecayLR",
+    "GuidedWaveringLR",
+    "FloorExponentialLR",
 ]
 
 import math
 from torch.optim import Optimizer
-from torch.optim.lr_scheduler import _LRScheduler
+from torch.optim.lr_scheduler import LRScheduler
+from typing import Optional
+from numbers import Number
+from lt_tensor.misc_utils import update_lr
 
 
-class WarmupDecayScheduler(_LRScheduler):
+class WarmupDecayScheduler(LRScheduler):
     def __init__(
         self,
         optimizer: Optimizer,
@@ -49,7 +54,7 @@ class WarmupDecayScheduler(_LRScheduler):
         return lrs
 
 
-class AdaptiveDropScheduler(_LRScheduler):
+class AdaptiveDropScheduler(LRScheduler):
     def __init__(
         self,
         optimizer,
@@ -89,26 +94,147 @@ class AdaptiveDropScheduler(_LRScheduler):
         return [group["lr"] for group in self.optimizer.param_groups]
 
 
-class WaveringLRScheduler(_LRScheduler):
+class SinusoidalDecayLR(LRScheduler):
     def __init__(
-        self, optimizer, base_lr, max_lr, period=1000, decay=0.999, last_epoch=-1
+        self,
+        optimizer: Optimizer,
+        initial_lr: float = 1e-3,
+        target_lr: float = 1e-5,
+        floor_lr: float = 1e-7,
+        decay_rate: float = 1e-6,  # decay per period
+        wave_amplitude: float = 1e-5,
+        period: int = 256,
+        last_epoch: int = -1,
     ):
-        """
-        Sinusoidal-like oscillating LR. Can escape shallow local minima.
-        - base_lr: minimum LR
-        - max_lr: maximum LR
-        - period: full sine cycle in steps
-        - decay: multiplies max_lr each cycle
-        """
-        self.base_lr = base_lr
-        self.max_lr = max_lr
+        assert decay_rate != 0.0, "decay_rate must be different from 0.0"
+        assert (
+            initial_lr >= target_lr >= floor_lr
+        ), "Must satisfy: initial_lr ≥ target_lr ≥ floor_lr"
+
+        self.initial_lr = initial_lr
+        self.target_lr = target_lr
+        self.floor_lr = floor_lr
+        self.decay_rate = decay_rate
+        self.wave_amplitude = wave_amplitude
         self.period = period
-        self.decay = decay
+
         super().__init__(optimizer, last_epoch)
 
     def get_lr(self):
-        cycle = self.last_epoch // self.period
-        step_in_cycle = self.last_epoch % self.period
-        factor = math.sin(math.pi * step_in_cycle / self.period)
-        amplitude = (self.max_lr - self.base_lr) * (self.decay**cycle)
-        return [self.base_lr + amplitude * factor for _ in self.optimizer.param_groups]
+        step = self.last_epoch + 1
+        cycles = step // self.period
+        t = step % self.period
+        # Decay center down to target_lr, then freeze
+        center_decay = math.exp(-self.decay_rate * cycles)
+        center = max(self.target_lr, self.initial_lr * center_decay)
+        # Decay amplitude in sync with center (relative to initial)
+        amplitude_decay = math.exp(-self.decay_rate * cycles)
+        current_amplitude = self.wave_amplitude * self.initial_lr * amplitude_decay
+        sin_offset = math.sin(2 * math.pi * t / self.period)
+        lr = max(center + current_amplitude * sin_offset, self.floor_lr)
+        return [lr for _ in self.optimizer.param_groups]
+
+
+class GuidedWaveringLR(LRScheduler):
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        initial_lr: float = 1e-3,
+        target_lr: float = 1e-5,
+        floor_lr: float = 1e-7,
+        decay_rate: float = 0.01,
+        wave_amplitude: float = 0.02,
+        period: int = 256,
+        stop_decay_after: int = None,
+        last_epoch: int = -1,
+    ):
+        assert decay_rate != 0.0, "decay_rate must be non-zero"
+        assert (
+            initial_lr >= target_lr >= floor_lr
+        ), "Must satisfy: initial ≥ target ≥ floor"
+
+        self.initial_lr = initial_lr
+        self.target_lr = target_lr
+        self.floor_lr = floor_lr
+        self.decay_rate = decay_rate
+        self.wave_amplitude = wave_amplitude
+        self.period = period
+        self.stop_decay_after = stop_decay_after
+
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        step = self.last_epoch + 1
+        cycles = step // self.period
+        t = step % self.period
+
+        decay_cycles = (
+            min(cycles, self.stop_decay_after) if self.stop_decay_after else cycles
+        )
+        center = max(
+            self.target_lr, self.initial_lr * math.exp(-self.decay_rate * decay_cycles)
+        )
+        amp = (
+            self.wave_amplitude
+            * self.initial_lr
+            * math.exp(-self.decay_rate * decay_cycles)
+        )
+        phase = 2 * math.pi * t / self.period
+        wave = math.sin(phase) * math.cos(phase)
+        lr = max(center + amp * wave, self.floor_lr)
+        return [lr for _ in self.optimizer.param_groups]
+
+
+class FloorExponentialLR(LRScheduler):
+    """Modified version from exponential lr, to have a minimum and reset functions.
+
+    Decays the learning rate of each parameter group by gamma every epoch.
+
+    When last_epoch=-1, sets initial lr as lr.
+
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        gamma (float): Multiplicative factor of learning rate decay.
+        last_epoch (int): The index of last epoch. Default: -1.
+    """
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        initial_lr: float = 1e-4,
+        gamma: float = 0.99998,
+        last_epoch: int = -1,
+        floor_lr: float = 1e-6,
+    ):
+        self.gamma = gamma
+        self.floor_lr = floor_lr
+        self.initial_lr = initial_lr
+
+        super().__init__(optimizer, last_epoch)
+
+    def set_floor(self, new_value: float):
+        assert isinstance(new_value, Number)
+        self.floor_lr = new_value
+
+    def reset_lr(self, new_value: Optional[float] = None):
+        new_lr = new_value if isinstance(new_value, Number) else self.initial_lr
+        self.initial_lr = new_lr
+        update_lr(self.optimizer, new_lr)
+
+    def get_lr(self):
+
+        if self.last_epoch == 0:
+            return [
+                max(group["lr"], self.floor_lr) for group in self.optimizer.param_groups
+            ]
+
+        return [
+            max(group["lr"] * self.gamma, self.floor_lr)
+            for group in self.optimizer.param_groups
+        ]
+
+    def _get_closed_form_lr(self):
+        return [
+            max(base_lr * self.gamma**self.last_epoch, self.floor_lr)
+            for base_lr in self.base_lrs
+        ]

lt_tensor/misc_utils.py

@@ -24,6 +24,7 @@ __all__ = [
     "plot_view",
     "get_weights",
     "get_activated_conv",
+    "update_lr",
 ]
 
 import re
@@ -77,6 +78,33 @@ def get_activated_conv(
     )
 
 
+def get_loss_average(losses: List[float]):
+    """A little helper for training, for example:
+    ```python
+    losses = []
+    for epoch in range(100):
+        for inp, label in dataloader:
+            optimizer.zero_grad()
+            out = model(inp)
+            loss = loss_fn(out, label)
+            optimizer.step()
+            losses.append(loss.item())
+        print(f"Epoch {epoch+1} | Loss: {get_loss_average(losses):.4f}")
+    """
+    if not losses:
+        return float("nan")
+    return sum(losses) / len(losses)
+
+
+def update_lr(optimizer: optim.Optimizer, new_value: float = 1e-4):
+    for param_group in optimizer.param_groups:
+        if isinstance(param_group["lr"], Tensor):
+            param_group["lr"].fill_(new_value)
+        else:
+            param_group["lr"] = new_value
+    return optimizer
+
+
 def plot_view(
     data: Dict[str, List[Any]],
     title: str = "Loss",
@@ -520,49 +548,14 @@ def sample_tensor(tensor: torch.Tensor, num_samples: int = 5):
     return flat[idx]
 
 
-class TorchCacheUtils:
-    cached_shortcuts: dict[str, Callable[[None], None]] = {}
-
-    has_cuda: bool = torch.cuda.is_available()
-    has_xpu: bool = torch.xpu.is_available()
-    has_mps: bool = torch.mps.is_available()
-
-    _ignore: list[str] = []
-
-    def __init__(self):
-        pass
-
-    def _apply_clear(self, device: str):
-        if device in self._ignore:
-            gc.collect()
-            return
-        try:
-            clear_fn = self.cached_shortcuts.get(
-                device, getattr(torch, device).empty_cache
-            )
-            if device not in self.cached_shortcuts:
-                self.cached_shortcuts.update({device: clear_fn})
-
-        except Exception as e:
-            print(e)
-            self._ignore.append(device)
-
-    def clear(self):
-        gc.collect()
-        if self.has_xpu:
-            self._apply_clear("xpu")
-        if self.has_cuda:
-            self._apply_clear("cuda")
-        if self.has_mps:
-            self._apply_clear("mps")
-        gc.collect()
-
-
-_clear_cache_cls = TorchCacheUtils()
-
-
 def clear_cache():
-    _clear_cache_cls.clear()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    if torch.mps.is_available():
+        torch.mps.empty_cache()
+    if torch.xpu.is_available():
+        torch.xpu.empty_cache()
+    gc.collect()
 
 
 @cache_wrapper

lt_tensor/model_zoo/activations/__init__.py

@@ -0,0 +1,3 @@
+from . import alias_free, snake
+
+__all__ = ["snake", "alias_free"]

lt_tensor/model_zoo/activations/alias_free/__init__.py

@@ -0,0 +1,3 @@
+from .act import *
+from .filter import *
+from .resample import *

lt_tensor/model_zoo/activations/{alias_free_torch → alias_free}/act.py

@@ -1,15 +1,17 @@
-import torch
 import torch.nn as nn
-import torch.nn.functional as F
-from .resample import UpSample1d, DownSample1d
-from .resample import UpSample2d, DownSample2d
+from lt_tensor.model_zoo.activations.alias_free.resample import (
+    UpSample2d,
+    DownSample2d,
+    UpSample1d,
+    DownSample1d,
+)
 
 
 class Activation1d(nn.Module):
 
     def __init__(
         self,
-        activation,
+        activation: nn.Module,
         up_ratio: int = 2,
         down_ratio: int = 2,
         up_kernel_size: int = 12,
@@ -34,7 +36,7 @@ class Activation2d(nn.Module):
 
     def __init__(
         self,
-        activation,
+        activation: nn.Module,
         up_ratio: int = 2,
         down_ratio: int = 2,
         up_kernel_size: int = 12,