lt-tensor 0.0.1a33__tar.gz → 0.0.1a35__tar.gz

{lt_tensor-0.0.1a33 → lt_tensor-0.0.1a35}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a33
+Version: 0.0.1a35
 Summary: General utilities for PyTorch and others. Built for general use.
 Home-page: https://github.com/gr1336/lt-tensor/
 Author: gr1336

{lt_tensor-0.0.1a33 → lt_tensor-0.0.1a35}/lt_tensor/__init__.py

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

lt_tensor-0.0.1a35/lt_tensor/losses.py

@@ -0,0 +1,281 @@
+__all__ = [
+    "masked_cross_entropy",
+    "adaptive_l1_loss",
+    "contrastive_loss",
+    "smooth_l1_loss",
+    "hybrid_loss",
+    "diff_loss",
+    "cosine_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)
+
+
+def adaptive_l1_loss(
+    inp: Tensor,
+    tgt: Tensor,
+    weight: Optional[Tensor] = None,
+    scale: float = 1.0,
+    inverted: bool = False,
+):
+
+    if weight is not None:
+        loss = torch.mean(torch.abs((inp - tgt) + weight.mean()))
+    else:
+        loss = torch.mean(torch.abs(inp - tgt))
+    loss *= scale
+    if inverted:
+        return -loss
+    return loss
+
+
+def smooth_l1_loss(inp: Tensor, tgt: Tensor, beta=1.0, weight=None):
+    diff = torch.abs(inp - tgt)
+    loss = torch.where(diff < beta, 0.5 * diff**2 / beta, diff - 0.5 * beta)
+    if weight is not None:
+        loss *= weight
+    return loss.mean()
+
+
+def contrastive_loss(x1: Tensor, x2: Tensor, label: Tensor, margin: float = 1.0):
+    # label == 1: similar, label == 0: dissimilar
+    dist = torch.nn.functional.pairwise_distance(x1, x2)
+    loss = label * dist**2 + (1 - label) * torch.clamp(margin - dist, min=0.0) ** 2
+    return loss.mean()
+
+
+def cosine_loss(inp, tgt):
+    cos = torch.nn.functional.cosine_similarity(inp, tgt, dim=-1)
+    return 1 - cos.mean()  # Lower is better
+
+
+def masked_cross_entropy(
+    logits: torch.Tensor,  # [B, T, V]
+    targets: torch.Tensor,  # [B, T]
+    lengths: torch.Tensor,  # [B]
+    reduction: str = "mean",
+) -> torch.Tensor:
+    """
+    CrossEntropyLoss with masking for variable-length sequences.
+    - logits: unnormalized scores [B, T, V]
+    - targets: ground truth indices [B, T]
+    - lengths: actual sequence lengths [B]
+    """
+    B, T, V = logits.size()
+    logits = logits.view(-1, V)
+    targets = targets.view(-1)
+
+    # Create mask
+    mask = torch.arange(T, device=lengths.device).expand(B, T) < lengths.unsqueeze(1)
+    mask = mask.reshape(-1)
+
+    # Apply CE only where mask == True
+    loss = F.cross_entropy(
+        logits[mask], targets[mask], reduction="mean" if reduction == "mean" else "none"
+    )
+    if reduction == "none":
+        return loss
+    return loss
+
+
+def diff_loss(pred_noise, true_noise, mask=None):
+    """Standard diffusion noise-prediction loss (e.g., DDPM)"""
+    if mask is not None:
+        return F.mse_loss(pred_noise * mask, true_noise * mask)
+    return F.mse_loss(pred_noise, true_noise)
+
+
+def hybrid_diff_loss(pred_noise, true_noise, alpha=0.5):
+    """Combines L1 and L2"""
+    l1 = F.l1_loss(pred_noise, true_noise)
+    l2 = F.mse_loss(pred_noise, true_noise)
+    return alpha * l1 + (1 - alpha) * l2
+
+
+def gan_d_loss(real_preds, fake_preds, use_lsgan=True):
+    loss = 0
+    for real, fake in zip(real_preds, fake_preds):
+        if use_lsgan:
+            loss += F.mse_loss(real, torch.ones_like(real)) + F.mse_loss(
+                fake, torch.zeros_like(fake)
+            )
+        else:
+            loss += -torch.mean(torch.log(real + 1e-7)) - torch.mean(
+                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-0.0.1a35/lt_tensor/lr_schedulers.py

@@ -0,0 +1,240 @@
+__all__ = [
+    "WarmupDecayScheduler",
+    "AdaptiveDropScheduler",
+    "SinusoidalDecayLR",
+    "GuidedWaveringLR",
+    "FloorExponentialLR",
+]
+
+import math
+from torch.optim import Optimizer
+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):
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        warmup_steps: int,
+        total_steps: int,
+        decay_type: str = "linear",  # or "cosine"
+        min_lr: float = 0.0,
+        last_epoch: int = -1,
+    ):
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.decay_type = decay_type
+        self.min_lr = min_lr
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        step = self.last_epoch + 1
+        warmup = self.warmup_steps
+        total = self.total_steps
+        lrs = []
+
+        for base_lr in self.base_lrs:
+            if step < warmup:
+                lr = base_lr * step / warmup
+            else:
+                progress = (step - warmup) / max(1, total - warmup)
+                if self.decay_type == "linear":
+                    lr = base_lr * (1.0 - progress)
+                elif self.decay_type == "cosine":
+                    lr = base_lr * 0.5 * (1 + math.cos(math.pi * progress))
+                else:
+                    raise ValueError(f"Unknown decay type: {self.decay_type}")
+
+            lr = max(self.min_lr, lr)
+            lrs.append(lr)
+
+        return lrs
+
+
+class AdaptiveDropScheduler(LRScheduler):
+    def __init__(
+        self,
+        optimizer,
+        drop_factor=0.5,
+        patience=10,
+        min_lr=1e-6,
+        cooldown=5,
+        last_epoch=-1,
+    ):
+        self.drop_factor = drop_factor
+        self.patience = patience
+        self.min_lr = min_lr
+        self.cooldown = cooldown
+        self.cooldown_counter = 0
+        self.best_loss = float("inf")
+        self.bad_steps = 0
+        super().__init__(optimizer, last_epoch)
+
+    def step(self, val_loss=None):
+        if val_loss is not None:
+            if val_loss < self.best_loss:
+                self.best_loss = val_loss
+                self.bad_steps = 0
+                self.cooldown_counter = 0
+            else:
+                self.bad_steps += 1
+                if self.bad_steps >= self.patience and self.cooldown_counter == 0:
+                    for i, group in enumerate(self.optimizer.param_groups):
+                        new_lr = max(group["lr"] * self.drop_factor, self.min_lr)
+                        group["lr"] = new_lr
+                    self.cooldown_counter = self.cooldown
+                    self.bad_steps = 0
+        if self.cooldown_counter > 0:
+            self.cooldown_counter -= 1
+
+    def get_lr(self):
+        return [group["lr"] for group in self.optimizer.param_groups]
+
+
+class SinusoidalDecayLR(LRScheduler):
+    def __init__(
+        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,
+    ):
+        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
+
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        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-0.0.1a33 → lt_tensor-0.0.1a35}/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-0.0.1a35/lt_tensor/model_zoo/activations/__init__.py

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

lt_tensor-0.0.1a35/lt_tensor/model_zoo/activations/alias_free/__init__.py

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

{lt_tensor-0.0.1a33/lt_tensor/model_zoo/activations/alias_free_torch → lt_tensor-0.0.1a35/lt_tensor/model_zoo/activations/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,