lt-tensor 0.0.1a9__py3-none-any.whl → 0.0.1a11__py3-none-any.whl

lt_tensor/losses.py

@@ -1,10 +1,106 @@
-__all__ = ["masked_cross_entropy"]
+__all__ = [
+    "masked_cross_entropy",
+    "adaptive_l1_loss",
+    "contrastive_loss",
+    "smooth_l1_loss",
+    "hybrid_loss",
+    "diff_loss",
+    "cosine_loss",
+    "gan_loss",
+    "ft_n_loss",
+]
 import math
 import random
 from .torch_commons import *
 from lt_utils.common import *
 import torch.nn.functional as F
 
+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
+
+
+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]
@@ -61,85 +157,3 @@ def gan_d_loss(real_preds, fake_preds, use_lsgan=True):
                 torch.log(1 - fake + 1e-7)
             )
     return loss
-
-
-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
-
-
-def gan_g_loss(fake_preds, use_lsgan=True):
-    loss = 0
-    for fake in fake_preds:
-        if use_lsgan:
-            loss += F.mse_loss(fake, torch.ones_like(fake))
-        else:
-            loss += -torch.mean(torch.log(fake + 1e-7))
-    return loss
-
-
-def feature_matching_loss(real_feats, fake_feats):
-    """real_feats and fake_feats are lists of intermediate features"""
-    loss = 0
-    for real_layers, fake_layers in zip(real_feats, fake_feats):
-        for r, f in zip(real_layers, fake_layers):
-            loss += F.l1_loss(f, r.detach())
-    return loss
-
-
-def feature_loss(real_fmaps, fake_fmaps, weight=2.0):
-    loss = 0.0
-    for dr, dg in zip(real_fmaps, fake_fmaps):  # Each (layer list from a discriminator)
-        for r_feat, g_feat in zip(dr, dg):
-            loss += F.l1_loss(r_feat, g_feat)
-    return loss * weight
-
-
-def discriminator_loss(disc_real_outputs, disc_generated_outputs):
-    loss = 0.0
-    r_losses = []
-    g_losses = []
-
-    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
-        r_loss = F.mse_loss(dr, torch.ones_like(dr))
-        g_loss = F.mse_loss(dg, torch.zeros_like(dg))
-        loss += r_loss + g_loss
-        r_losses.append(r_loss)
-        g_losses.append(g_loss)
-
-    return loss, r_losses, g_losses
-
-
-def generator_loss(fake_outputs):
-    total = 0.0
-    g_losses = []
-    for out in fake_outputs:
-        loss = F.mse_loss(out, torch.ones_like(out))
-        g_losses.append(loss)
-        total += loss
-    return total, g_losses
-
-
-def multi_resolution_stft_loss(y, y_hat, fft_sizes=[512, 1024, 2048]):
-    loss = 0
-    for fft_size in fft_sizes:
-        hop = fft_size // 4
-        win = fft_size
-        y_stft = torch.stft(
-            y, n_fft=fft_size, hop_length=hop, win_length=win, return_complex=True
-        )
-        y_hat_stft = torch.stft(
-            y_hat, n_fft=fft_size, hop_length=hop, win_length=win, return_complex=True
-        )
-
-        loss += F.l1_loss(torch.abs(y_stft), torch.abs(y_hat_stft))
-    return loss

lt_tensor/misc_utils.py

@@ -15,6 +15,7 @@ __all__ = [
     "TorchCacheUtils",
     "clear_cache",
     "default_device",
+    "soft_restore",
     "Packing",
     "Padding",
     "Masking",
@@ -35,14 +36,29 @@ from lt_utils.misc_utils import ff_list
 import torch.nn.functional as F
 
 
+def soft_restore(tensor, epsilon=1e-6):
+    return torch.where(tensor == 0, torch.full_like(tensor, epsilon), tensor)
+
+
 def try_torch(fn: str, *args, **kwargs):
+    tryed_torch = False
+    not_present_message = (
+        f"Both `torch` and `torch.nn.functional` does not contain the module `{fn}`"
+    )
     try:
-        return getattr(F, fn)(*args, **kwargs)
-    except Exception as e:
-        try:
+        if hasattr(F, fn):
+            return getattr(F, fn)(*args, **kwargs)
+        elif hasattr(torch, fn):
+            tryed_torch = True
             return getattr(torch, fn)(*args, **kwargs)
+        return not_present_message
+    except Exception as a:
+        try:
+            if not tryed_torch and hasattr(torch, fn):
+                return getattr(torch, fn)(*args, **kwargs)
+            return str(a)
         except Exception as e:
-            return str(e)
+            return str(e) + " | " + str(a)
 
 
 def log_tensor(
@@ -152,7 +168,12 @@ class LogTensor:
             print(message)
             sys.stdout.flush()
 
-    def _process(self, name: str, resp: Union[Callable[[Any], Any], Any]):
+    def _process(
+        self,
+        name: str,
+        resp: Union[Callable[[Any], Any], Any],
+        do_not_print: bool = False,
+    ):
         if callable(resp):
             try:
                 response = resp()
@@ -161,15 +182,16 @@ class LogTensor:
                 print(e)
         else:
             response = resp
-        msg = self._setup_message(name, response)
-        self._print(msg)
-        return dict(name=name, results=response, message=msg)
+        if not do_not_print and self.do_print:
+            msg = self._setup_message(name, response)
+            self._print(msg)
+        return dict(item=name, value=response)
 
     def __call__(
         self,
-        main_item: Union[Tensor, np.ndarray, Sequence[Number]],
+        inputs: Union[Tensor, np.ndarray, Sequence[Number]],
         title: Optional[str] = None,
-        loss_targets: Optional[Union[Tensor, np.ndarray, Sequence[Number]]] = None,
+        target: Optional[Union[Tensor, np.ndarray, Sequence[Number]]] = None,
         *,
         log_tensor: bool = False,
         log_device: bool = False,
@@ -177,7 +199,7 @@ class LogTensor:
         log_std: bool = False,
         dim_mean: int = -1,
         dim_std: int = -1,
-        do_print_external: bool = False,
+        print_extended: bool = False,
         external_logs: List[Tuple[str, Union[Dict[str, Any], List[Any], Any]]] = [
             ("softmax", {"dim": 0}),
             ("relu", None),
@@ -186,12 +208,11 @@ class LogTensor:
             ("max", dict(dim=-1)),
         ],
         validate_item_type: bool = False,
+        exclude_invalid_losses: bool = True,
         **kwargs,
     ):
-        if kwargs:
-            self.update_settings(**kwargs)
-        invalid_type = not isinstance(main_item, (Tensor, np.ndarray, list, tuple))
-        _main_item_tp = type(main_item)
+        invalid_type = not isinstance(inputs, (Tensor, np.ndarray, list, tuple))
+        _main_item_tp = type(inputs)
         assert (
             not validate_item_type or not invalid_type
         ), f"Invalid Type: {_main_item_tp}"
@@ -199,8 +220,8 @@ class LogTensor:
             self._print(f"Invalid Type: {_main_item_tp}")
             return
 
-        main_item = self._setup_tensor(main_item)
-        loss_targets = self._setup_tensor(loss_targets)
+        inputs = self._setup_tensor(inputs)
+        target = self._setup_tensor(target)
         if is_str(title):
             title = re.sub(r"\s+", " ", title.replace("_", " "))
             has_title = is_str(title)
@@ -209,43 +230,43 @@ class LogTensor:
         else:
             title = "Unnamed"
 
-        current_register = {"title": title}
-
-        current_register.update(self._process("shape", main_item.shape))
-        current_register.update(self._process("ndim", main_item.ndim))
-        current_register.update(self._process("dtype", main_item.dtype))
+        current_register = {"title": title, "values": []}
+        current_register["shape"] = self._process("shape", inputs.shape)
+        current_register["ndim"] = self._process("ndim", inputs.ndim)
+        current_register["dtype"] = self._process("dtype", inputs.dtype)
         if log_device:
-            current_register.update(self._process("device", main_item.device))
+            current_register["device"] = self._process("device", inputs.device)
 
         if log_mean:
-            fn = lambda: main_item.mean(
+            fn = lambda: inputs.mean(
                 dim=dim_mean,
             )
-            current_register.update(self._process("mean", fn))
+            current_register["mean"] = self._process("mean", fn)
 
         if log_std:
-            fn = lambda: main_item.std(dim=dim_std)
-            current_register.update(self._process("std", fn))
+            fn = lambda: inputs.std(dim=dim_std)
+            current_register["std"] = self._process("std", fn)
 
         if external_logs:
             old_print = self.do_print
-            self.do_print = do_print_external
+            self.do_print = print_extended
             self._print("\n---[ External Logs ] ---")
             for log_fn, log_args in external_logs:
                 if isinstance(log_args, Sequence) and not isinstance(log_args, str):
-                    value = try_torch(log_fn, main_item, *log_args)
+                    value = try_torch(log_fn, inputs, *log_args)
                 elif isinstance(log_args, dict):
-                    value = try_torch(log_fn, main_item, **log_args)
+                    value = try_torch(log_fn, inputs, **log_args)
                 elif log_args is None:
-                    value = try_torch(log_fn, main_item)
+                    value = try_torch(log_fn, inputs)
                 else:
-                    value = try_torch(log_fn, main_item, log_args)
+                    value = try_torch(log_fn, inputs, log_args)
+
                 results = self._process(log_fn, value)
                 current_register[log_fn] = results
             self.do_print = old_print
 
-        if loss_targets is not None:
-            losses = get_losses(main_item, loss_targets, False)
+        if target is not None:
+            losses = get_losses(inputs, target, exclude_invalid_losses)
             started_ls = False
             if self.do_print:
                 for loss, res in losses.items():
@@ -258,7 +279,9 @@ class LogTensor:
             current_register["loss"] = losses
 
         if log_tensor:
-            current_register.update(self._process("Tensor", main_item))
+            current_register["values"].append(
+                self._process("Tensor", inputs, not print_extended)
+            )
 
         self._print(self.end_with)
         self._store_item_and_update(current_register)

lt_tensor/model_base.py

@@ -1,10 +1,12 @@
-__all__ = ["Model"]
-
+__all__ = ["Model", "_ModelExtended", "LossTracker"]
 
+import gc
+import json
+import math
 import warnings
 from .torch_commons import *
 from lt_utils.common import *
-from lt_utils.misc_utils import log_traceback
+from lt_utils.misc_utils import log_traceback, get_current_time
 
 T = TypeVar("T")
 
@@ -17,6 +19,80 @@ POSSIBLE_OUTPUT_TYPES: TypeAlias = Union[
 ]
 
 
+class LossTracker:
+    last_file = f"logs/history_{get_current_time()}.json"
+
+    def __init__(self, max_len=50_000):
+        self.max_len = max_len
+        self.history = {
+            "train": [],
+            "eval": [],
+        }
+
+    def append(self, loss: float, mode: Literal["train", "eval"] = "train"):
+        assert mode in self.history, f"Invalid mode '{mode}'. Use 'train' or 'eval'."
+        self.history[mode].append(float(loss))
+        if len(self.history[mode]) > self.max_len:
+            self.history[mode] = self.history[mode][-self.max_len :]
+
+    def get(self, mode: Literal["train", "eval"] = "train"):
+        return self.history.get(mode, [])
+
+    def save(self, path: Optional[PathLike] = None):
+        if path is None:
+            path = f"logs/history_{get_current_time()}.json"
+
+        Path(path).parent.mkdir(exist_ok=True, parents=True)
+        with open(path, "w") as f:
+            json.dump(self.history, f, indent=2)
+
+        self.last_file = path
+
+    def load(self, path: Optional[PathLike] = None):
+        if path is None:
+            _path = self.last_file
+        else:
+            _path = path
+        with open(_path) as f:
+            self.history = json.load(f)
+        if path is not None:
+            self.last_file = path
+
+    def plot(self, backend: Literal["matplotlib", "plotly"] = "plotly"):
+        if backend == "plotly":
+            try:
+                import plotly.graph_objs as go
+            except ModuleNotFoundError:
+                warnings.warn(
+                    "No installation of plotly was found. To use it use 'pip install plotly' and restart this application!"
+                )
+                return
+            fig = go.Figure()
+            for mode, losses in self.history.items():
+                if losses:
+                    fig.add_trace(go.Scatter(y=losses, name=mode.capitalize()))
+            fig.update_layout(
+                title="Training vs Evaluation Loss",
+                xaxis_title="Step",
+                yaxis_title="Loss",
+                template="plotly_dark",
+            )
+            fig.show()
+
+        elif backend == "matplotlib":
+            import matplotlib.pyplot as plt
+
+            for mode, losses in self.history.items():
+                if losses:
+                    plt.plot(losses, label=f"{mode.capitalize()} Loss")
+            plt.title("Loss over Time")
+            plt.xlabel("Step")
+            plt.ylabel("Loss")
+            plt.legend()
+            plt.grid(True)
+            plt.show()
+
+
 class Model(nn.Module, ABC):
     """
     This makes it easier to assign a device and retrieves it later
@@ -24,6 +100,8 @@ class Model(nn.Module, ABC):
 
     _device: torch.device = ROOT_DEVICE
     _autocast: bool = False
+    _loss_history: LossTracker = LossTracker(100_000)
+    _is_unfrozen: bool = False
 
     @property
     def autocast(self):
@@ -61,6 +139,7 @@ class Model(nn.Module, ABC):
             if hasattr(self, weight):
                 w = getattr(self, weight)
                 if isinstance(w, nn.Module):
+
                     w.requires_grad_(not freeze)
         else:
             weight.requires_grad_(not freeze)
@@ -112,21 +191,27 @@ class Model(nn.Module, ABC):
         for name, param in self.named_parameters():
             if no_exclusions:
                 try:
-                    param.requires_grad_(False)
-                    frozen.append(name)
+                    if param.requires_grad:
+                        param.requires_grad_(False)
+                        frozen.append(name)
+                    else:
+                        not_frozen.append((name, "was_frozen"))
                 except Exception as e:
                     not_frozen.append((name, str(e)))
             elif any(layer in name for layer in exclude):
                 try:
-                    param.requires_grad_(False)
-                    frozen.append(name)
+                    if param.requires_grad:
+                        param.requires_grad_(False)
+                        frozen.append(name)
+                    else:
+                        not_frozen.append((name, "was_frozen"))
                 except Exception as e:
                     not_frozen.append((name, str(e)))
             else:
-                not_frozen.append((name, "Excluded"))
+                not_frozen.append((name, "excluded"))
         return dict(frozen=frozen, not_frozen=not_frozen)
 
-    def unfreeze_all_except(self, exclude: Optional[list[str]] = None):
+    def unfreeze_all(self, exclude: Optional[list[str]] = None):
         """Unfreezes all model parameters except specified layers."""
         no_exclusions = not exclude
         unfrozen = []
@@ -134,18 +219,24 @@ class Model(nn.Module, ABC):
         for name, param in self.named_parameters():
             if no_exclusions:
                 try:
-                    param.requires_grad_(True)
-                    unfrozen.append(name)
+                    if not param.requires_grad:
+                        param.requires_grad_(True)
+                        unfrozen.append(name)
+                    else:
+                        not_unfrozen.append((name, "was_unfrozen"))
                 except Exception as e:
                     not_unfrozen.append((name, str(e)))
             elif any(layer in name for layer in exclude):
                 try:
-                    param.requires_grad_(True)
-                    unfrozen.append(name)
+                    if not param.requires_grad:
+                        param.requires_grad_(True)
+                        unfrozen.append(name)
+                    else:
+                        not_unfrozen.append((name, "was_unfrozen"))
                 except Exception as e:
                     not_unfrozen.append((name, str(e)))
             else:
-                not_unfrozen.append((name, "Excluded"))
+                not_unfrozen.append((name, "excluded"))
         return dict(unfrozen=unfrozen, not_unfrozen=not_unfrozen)
 
     def to(self, *args, **kwargs):
@@ -192,6 +283,7 @@ class Model(nn.Module, ABC):
 
         self._apply(convert)
         self.device = device
+        self._apply_device_to()
         return self
 
     def ipu(self, device: Optional[Union[int, torch.device]] = None) -> T:
@@ -202,6 +294,7 @@ class Model(nn.Module, ABC):
                 ":" + str(device) if isinstance(device, (int, float)) else device.index
             )
         self.device = dvc
+        self._apply_device_to()
         return self
 
     def xpu(self, device: Optional[Union[int, torch.device]] = None) -> T:
@@ -212,6 +305,7 @@ class Model(nn.Module, ABC):
                 ":" + str(device) if isinstance(device, (int, float)) else device.index
             )
         self.device = dvc
+        self._apply_device_to()
         return self
 
     def cuda(self, device: Optional[Union[int, torch.device]] = None) -> T:
@@ -222,6 +316,7 @@ class Model(nn.Module, ABC):
                 ":" + str(device) if isinstance(device, (int, float)) else device.index
             )
         self.device = dvc
+        self._apply_device_to()
         return self
 
     def mtia(self, device: Optional[Union[int, torch.device]] = None) -> T:
@@ -232,11 +327,13 @@ class Model(nn.Module, ABC):
                 ":" + str(device) if isinstance(device, (int, float)) else device.index
             )
         self.device = dvc
+        self._apply_device_to()
         return self
 
     def cpu(self) -> T:
         super().cpu()
         self.device = "cpu"
+        self._apply_device_to()
         return self
 
     def count_trainable_parameters(self, module_name: Optional[str] = None):
@@ -314,15 +411,26 @@ class Model(nn.Module, ABC):
         else:
             print(f"Non-Trainable Parameters: {params}")
 
-    def save_weights(self, path: Union[Path, str]):
+    def save_weights(
+        self,
+        path: Union[Path, str],
+        replace: bool = False,
+    ):
         path = Path(path)
+        model_dir = path
         if path.exists():
-            assert (
-                path.is_file()
-            ), "The provided path exists but its a directory not a file!"
-            path.rmdir()
+            if path.is_dir():
+                model_dir = Path(path, f"model_{get_current_time()}.pt")
+            elif path.is_file():
+                if replace:
+                    path.unlink()
+                else:
+                    model_dir = Path(path.parent, f"model_{get_current_time()}.pt")
+        else:
+            if not "." in str(path):
+                model_dir = Path(path, f"model_{get_current_time()}.pt")
         path.parent.mkdir(exist_ok=True, parents=True)
-        torch.save(obj=self.state_dict(), f=str(path))
+        torch.save(obj=self.state_dict(), f=str(model_dir))
 
     def load_weights(
         self,
@@ -338,7 +446,14 @@ class Model(nn.Module, ABC):
         if not path.exists():
             assert not raise_if_not_exists, "Path does not exists!"
             return None
-        assert path.is_file(), "The provided path is not a valid file!"
+        if path.is_dir():
+            possible_files = list(Path(path).rglob("*.pt"))
+            assert (
+                possible_files or not raise_if_not_exists
+            ), "No model could be found in the given path!"
+            if not possible_files:
+                return None
+            path = sorted(possible_files)[-1]
         state_dict = torch.load(
             str(path), weights_only=weights_only, mmap=mmap, **torch_loader_kwargs
         )
@@ -353,30 +468,131 @@ class Model(nn.Module, ABC):
     def inference(self, *args, **kwargs):
         if self.training:
             self.eval()
-        if self.autocast:
-            with torch.autocast(device_type=self.device.type):
-                return self(*args, **kwargs)
         return self(*args, **kwargs)
 
     def train_step(
         self,
-        *args,
+        *inputs,
         **kwargs,
     ):
         """Train Step"""
         if not self.training:
             self.train()
-        return self(*args, **kwargs)
-
-    @torch.autocast(device_type=_device.type)
-    def ac_forward(self, *args, **kwargs):
-        return
+        return self(*inputs, **kwargs)
 
     def __call__(self, *args, **kwds) -> POSSIBLE_OUTPUT_TYPES:
-        return super().__call__(*args, **kwds)
+        if self.autocast and not self.training:
+            with torch.autocast(device_type=self.device.type):
+                return super().__call__(*args, **kwds)
+        else:
+            return super().__call__(*args, **kwds)
 
     @abstractmethod
     def forward(
         self, *args, **kwargs
     ) -> Union[Tensor, Sequence[Tensor], Dict[Any, Union[Any, Tensor]]]:
         pass
+
+    def add_loss(
+        self, loss: Union[float, list[float]], mode: Literal["train", "eval"] = "train"
+    ):
+        if isinstance(loss, Number) and loss:
+            self._loss_history.append(loss, mode)
+        elif isinstance(loss, (list, tuple)):
+            if loss:
+                self._loss_history.append(sum(loss) / len(loss), mode=mode)
+        elif isinstance(loss, Tensor):
+            try:
+                self._loss_history.append(loss.detach().flatten().mean().item())
+            except Exception as e:
+                log_traceback(e, "add_loss - Tensor")
+
+    def save_loss_history(self, path: Optional[PathLike] = None):
+        self._loss_history.save(path)
+
+    def load_loss_history(self, path: Optional[PathLike] = None):
+        self._loss_history.load(path)
+
+    def get_loss_avg(self, mode: Literal["train", "eval"], quantity: int = 0):
+        t_list = self._loss_history.get("train")
+        if not t_list:
+            return float("nan")
+        if quantity > 0:
+            t_list = t_list[-quantity:]
+        return sum(t_list) / len(t_list)
+
+    def freeze_unfreeze_loss(
+        self,
+        losses: Optional[Union[float, List[float]]] = None,
+        trigger_loss: float = 0.1,
+        excluded_modules: Optional[List[str]] = None,
+        eval_last: int = 1000,
+    ):
+        """If a certain threshold is reached the weights will freeze or unfreeze the modules.
+        the biggest use-case for this function is when training GANs where the balance
+        from the discriminator and generator must be kept.
+
+        Args:
+            losses (Union[float, List[float]], Optional): The loss value or a list of losses that will be used to determine if it has reached or not the threshold. Defaults to None.
+            trigger_loss (float, optional): The value where the weights will be either freeze or unfreeze. Defaults to 0.1.
+            excluded_modules (list[str], optional): The list of modules (names) that is not to be changed by either freezing nor unfreezing. Defaults to None.
+            eval_last (float, optional): The number of previous losses to be locked behind to calculate the current averange. Default to 1000.
+
+        returns:
+            bool: True when its frozen and false when its trainable.
+        """
+        if losses is not None:
+            calculated = None
+            self.add_loss(losses)
+
+        value = self.get_loss_avg("train", eval_last)
+
+        if value <= trigger_loss:
+            if self._is_unfrozen:
+                self.freeze_all(excluded_modules)
+                self._is_unfrozen = False
+            return True
+        else:
+            if not self._is_unfrozen:
+                self.unfreeze_all(excluded_modules)
+                self._is_unfrozen = True
+            return False
+
+
+class _ModelExtended(Model):
+    """Planed, but not ready, maybe in the near future?"""
+    criterion: Optional[Callable[[Tensor, Tensor], Tensor]] = None
+    optimizer: Optional[optim.Optimizer] = None
+
+    def train_step(
+        self,
+        *inputs,
+        loss_label: Optional[Tensor] = None,
+        **kwargs,
+    ):
+        if not self.training:
+            self.train()
+        if self.optimizer is not None:
+            self.optimizer.zero_grad()
+        if self.autocast:
+            if self.criterion is None:
+                raise RuntimeError(
+                    "To use autocast during training, you must assign a criterion first!"
+                )
+            with torch.autocast(device_type=self.device.type):
+                out = self.forward(*loss_label, **kwargs)
+                loss = self.criterion(out, loss_label)
+
+            if self.optimizer is not None:
+                loss.backward()
+                self.optimizer.step()
+            return loss
+        elif self.criterion is not None:
+            out = self.forward(*loss_label, **kwargs)
+            loss = self.criterion(out, loss_label)
+            if self.optimizer is not None:
+                loss.backward()
+                self.optimizer.step()
+            return loss
+        else:
+            return self(*inputs, **kwargs)

lt_tensor/model_zoo/bsc.py

@@ -208,3 +208,25 @@ class MultiScaleEncoder1D(Model):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # x: [B, C, T]
         return self.net(x)  # [B, hidden, T]
+    
+    
+class AudioClassifier(Model):
+    def __init__(self, n_mels:int=80, num_classes=5):
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.Conv1d(n_mels, 256, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(256, 256, kernel_size=3, padding=1, groups=4),
+            nn.BatchNorm1d(256),
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(256, 256, kernel_size=3, padding=1),
+            nn.BatchNorm1d(256),
+            nn.LeakyReLU(0.2),
+            nn.AdaptiveAvgPool1d(1),  # Output shape: [B, 64, 1]
+            nn.Flatten(),  # -> [B, 64]
+            nn.Linear(256, num_classes),
+        )
+        self.eval()
+
+    def forward(self, x):
+        return self.model(x)

lt_tensor/model_zoo/disc.py

@@ -76,20 +76,6 @@ class PeriodDiscriminator(Model):
         return x.flatten(1, -1), f_map
 
 
-class MultiPeriodDiscriminator(Model):
-    def __init__(self, periods=[2, 3, 5, 7, 11]):
-        super().__init__()
-
-        self.discriminators = nn.ModuleList([PeriodDiscriminator(p) for p in periods])
-
-    def forward(self, x: torch.Tensor):
-        """
-        x: (B, T)
-        Returns: list of tuples of outputs from each period discriminator and the f_map.
-        """
-        return [d(x) for d in self.discriminators]
-
-
 class ScaleDiscriminator(nn.Module):
     def __init__(self, use_spectral_norm=False):
         super().__init__()
@@ -123,11 +109,11 @@ class ScaleDiscriminator(nn.Module):
 
 
 class MultiScaleDiscriminator(Model):
-    def __init__(self):
+    def __init__(self, layers: int = 3):
         super().__init__()
         self.pooling = nn.AvgPool1d(4, 2, padding=2)
         self.discriminators = nn.ModuleList(
-            [ScaleDiscriminator(i == 0) for i in range(3)]
+            [ScaleDiscriminator(i == 0) for i in range(layers)]
         )
 
     def forward(self, x: torch.Tensor):
@@ -136,57 +122,75 @@ class MultiScaleDiscriminator(Model):
         Returns: list of outputs from each scale discriminator
         """
         outputs = []
+        features = []
         for i, d in enumerate(self.discriminators):
             if i != 0:
                 x = self.pooling(x)
-            outputs.append(d(x))
-        return outputs
+            out, f_map = d(x)
+            outputs.append(out)
+            features.append(f_map)
+        return outputs, features
 
 
-class GeneralLossDescriminator(Model):
-    """TODO: build an unified loss for both mpd and msd here."""
-
-    def __init__(self):
+class MultiPeriodDiscriminator(Model):
+    def __init__(self, periods: List[int] = [2, 3, 5, 7, 11]):
         super().__init__()
-        self.mpd = MultiPeriodDiscriminator()
-        self.msd = MultiScaleDiscriminator()
-        self.print_trainable_parameters()
-
-    def _get_group_(self):
-        pass
+        self.discriminators = nn.ModuleList([PeriodDiscriminator(p) for p in periods])
 
-    def forward(self, x: Tensor, y_hat: Tensor):
-        return
+    def forward(self, x: torch.Tensor):
+        """
+        x: (B, T)
+        Returns: list of tuples of outputs from each period discriminator and the f_map.
+        """
+        # torch.log(torch.clip(x, min=clip_val))
+        out_map = []
+        feat_map = []
+        for d in self.discriminators:
+            out, feat = d(x)
+            out_map.append(out)
+            feat_map.append(feat)
+        return out_map, feat_map
 
 
-def discriminator_loss(d_outputs_real, d_outputs_fake):
+def discriminator_loss(real_out_map, fake_out_map):
     loss = 0.0
-    for real_out, fake_out in zip(d_outputs_real, d_outputs_fake):
-        real_score = real_out[0]
-        fake_score = fake_out[0]
-        loss += torch.mean(F.relu(1.0 - real_score)) + torch.mean(
-            F.relu(1.0 + fake_score)
-        )
-    return loss
+    rl, fl = [], []
+    for real_out, fake_out in zip(real_out_map, fake_out_map):
+        real_loss = torch.mean((1.0 - real_out) ** 2)
+        fake_loss = torch.mean(fake_out**2)
+        loss += real_loss + fake_loss
+        rl.append(real_loss.item())
+        fl.append(fake_loss.item())
+    return loss, sum(rl), sum(fl)
 
 
-def generator_adv_loss(d_outputs_fake):
+def generator_adv_loss(fake_disc_outputs: List[Tensor]):
     loss = 0.0
-    for fake_out in d_outputs_fake:
+    for fake_out in fake_disc_outputs:
         fake_score = fake_out[0]
         loss += -torch.mean(fake_score)
     return loss
 
 
-def feature_matching_loss(
-    d_outputs_real,
-    d_outputs_fake,
-    loss_fn: Callable[[Tensor, Tensor], Tensor] = F.mse_loss,
+def feature_loss(
+    fmap_r,
+    fmap_g,
+    weight=2.0,
+    loss_fn: Callable[[Tensor, Tensor], Tensor] = F.l1_loss,
 ):
     loss = 0.0
-    for real_out, fake_out in zip(d_outputs_real, d_outputs_fake):
-        real_feats = real_out[1]
-        fake_feats = fake_out[1]
-        for real_f, fake_f in zip(real_feats, fake_feats):
-            loss += loss_fn(fake_f, real_f)
-    return loss
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss += loss_fn(rl - gl)
+    return loss * weight
+
+
+def generator_loss(disc_generated_outputs):
+    loss = 0.0
+    gen_losses = []
+    for dg in disc_generated_outputs:
+        l = torch.mean((1.0 - dg) ** 2)
+        gen_losses.append(l.item())
+        loss += l
+
+    return loss, gen_losses

lt_tensor/model_zoo/istft.py

@@ -106,44 +106,3 @@ class Generator(Model):
         classname = m.__class__.__name__
         if "Conv" in classname:
             m.weight.data.normal_(mean, std)
-
-
-# Below are items found in the Rishikesh's repo that might work for this generator.
-# https://github.com/rishikksh20/iSTFTNet-pytorch/blob/781480e9563d4dff5a8cc9ef1af6c6e0cab025c8/models.py
-
-
-def feature_loss(fmap_r, fmap_g, weight=2.0):
-    """Feature matching loss between real and generated feature maps."""
-    loss = 0.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 * weight
-
-
-def discriminator_loss(disc_real_outputs, disc_generated_outputs):
-    """LSGAN-style loss for real and fake predictions."""
-    loss = 0.0
-    r_losses, g_losses = [], []
-
-    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
-        r_loss = torch.mean((1.0 - dr) ** 2)
-        g_loss = torch.mean(dg**2)
-        loss += r_loss + g_loss
-        r_losses.append(r_loss.item())
-        g_losses.append(g_loss.item())
-
-    return loss, r_losses, g_losses
-
-
-def generator_loss(disc_generated_outputs):
-    """LSGAN generator loss encouraging fake to look like real (close to 1)."""
-    loss = 0.0
-    gen_losses = []
-
-    for dg in disc_generated_outputs:
-        l = torch.mean((1.0 - dg) ** 2)
-        gen_losses.append(l.item())
-        loss += l
-
-    return loss, gen_losses

lt_tensor/noise_tools.py

@@ -271,9 +271,8 @@ class NoiseSchedulerB(nn.Module):
     def forward(
         self, x_0: Tensor, t: int, noise: Optional[Union[Tensor, float]] = None
     ) -> Tensor:
-        apply_noise()
         assert (
-            0 >= t < self.timesteps
+             0 <= t < self.timesteps
         ), f"Time step t={t} is out of bounds for scheduler with {self.timesteps} steps."
 
         if noise is None:

lt_tensor/transform.py

@@ -420,44 +420,11 @@ class InverseTransform(Model):
         self.onesided = onesided
         self.normalized = normalized
         self.window = torch.hann_window(win_length) if window is None else window
-        self.update_settings()
-        
+
     def _apply_device_to(self):
         """Applies to device while used with module `Model`"""
         self.window = self.window.to(device=self.device)
 
-    def update_settings(
-        self,
-        *,
-        n_fft: Optional[int] = None,
-        hop_length: Optional[int] = None,
-        win_length: Optional[int] = None,
-        length: Optional[int] = None,
-        window: Optional[Tensor] = None,
-        onesided: Optional[bool] = None,
-        return_complex: Optional[bool] = None,
-        center: Optional[bool] = None,
-        normalized: Optional[bool] = None,
-        **_,
-    ):
-
-        self.kwargs = dict(
-            n_fft=default(n_fft, self.n_fft),
-            hop_length=default(hop_length, self.hop_length),
-            win_length=default(win_length, self.win_length),
-            length=default(length, self.length),
-            window=default(window, self.window),
-            onesided=default(onesided, self.onesided),
-            return_complex=default(return_complex, self.return_complex),
-            center=default(center, self.center),
-            normalized=default(normalized, self.normalized),
-        )
-        if self.kwargs["onesided"] and self.kwargs["return_complex"]:
-            warnings.warn(
-                "You cannot use return_complex with `onesided` enabled. `return_complex` is set to False."
-            )
-            self.kwargs["return_complex"] = False
-
     def forward(self, spec: Tensor, phase: Tensor, **kwargs):
         """
         Perform the inverse short-time Fourier transform.
@@ -476,7 +443,16 @@ class InverseTransform(Model):
         Tensor
             Time-domain waveform reconstructed from `spec` and `phase`.
         """
-        if kwargs:
-            self.update_settings(**kwargs)
 
-        return torch.istft(spec * torch.exp(phase * 1j), **self.kwargs)
+        return torch.istft(
+            spec * torch.exp(phase * 1j),
+            n_fft = self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=self.window,
+            center=self.center,
+            normalized=self.normalized,
+            onesided=self.onesided,
+            length=self.length,
+            return_complex=self.return_complex,
+        )

{lt_tensor-0.0.1a9.dist-info → lt_tensor-0.0.1a11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a9
+Version: 0.0.1a11
 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.1a9.dist-info → lt_tensor-0.0.1a11.dist-info}/RECORD

@@ -1,28 +1,28 @@
 lt_tensor/__init__.py,sha256=uwJ7uiO18VYj8Z1V4KSOQ3ZrnowSgJWKCIiFBrzLMOI,429
-lt_tensor/losses.py,sha256=TinZJP2ypZ7Tdg6d9nnFWFkPyormfgQ0Z9P2ER3sqzE,4341
+lt_tensor/losses.py,sha256=1wrke1e68hUBNAoPdJgKni0pJvXKcieza_R8nwBzMW4,4937
 lt_tensor/lr_schedulers.py,sha256=LSZzqrOOLzSthD8k-W4cYPJt0vCjmHkiJkLr5e3yRTE,3659
 lt_tensor/math_ops.py,sha256=ewIYkvxIy_Lab_9ExjFUgLs-oYLOu8IRRDo7f1pn3i8,2248
-lt_tensor/misc_utils.py,sha256=-MGqs3Dt35M7EsgBqqTLin9v0xVu8V29HQZVtsxK5QU,24746
-lt_tensor/model_base.py,sha256=8qN7oklALFanOz-eqVzdnB9RD2kN_3ltynSMAPOl-TI,13413
+lt_tensor/misc_utils.py,sha256=8LqtpmLKqCo79NdH160ByQojG8YTDcw8aHKFgOFGVLI,25425
+lt_tensor/model_base.py,sha256=a2ogixC2fUyOLqz15TzCRcGXvBam--TdmpG83jw9Of8,21543
 lt_tensor/monotonic_align.py,sha256=LhBd8p1xdBzg6jQrQX1j7b4PNeYGwIqM24zcU-pHOLE,2239
-lt_tensor/noise_tools.py,sha256=JkWw0-bCMRNNMShwXKKt5KbO3104tvNiBePt-ThPkEo,11366
+lt_tensor/noise_tools.py,sha256=rfFbPsrsycWVuH9G4zZCQC9Vgi9r8hDaECcB0TZYSYQ,11345
 lt_tensor/torch_commons.py,sha256=fntsEU8lhBQo0ebonI1iXBkMbWMN3HpBsG13EWlP5s8,718
-lt_tensor/transform.py,sha256=Bxh87vFRKuZay_g1Alf_ZtEo89CzmV3XUQDINwHB7iA,14505
+lt_tensor/transform.py,sha256=LZZ9G7ud1cojERC7N7hMAbH9GC3ImY1hBIY00kVMs-I,13492
 lt_tensor/datasets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 lt_tensor/datasets/audio.py,sha256=YREyRsCvy-KS5tE0JNMWEdlIJogE1khLqhiq4wOWXVg,3777
 lt_tensor/model_zoo/__init__.py,sha256=jipEk50_DTMQbGg8FnDDukxmh7Bcwvl_QVRS3rkb7aY,283
-lt_tensor/model_zoo/bsc.py,sha256=muxIR7dU-Pvf-HFE-iy3zmRb1sTJlcs1vqdlnbU1Hss,6307
-lt_tensor/model_zoo/disc.py,sha256=SphFVVPZLP96-mZPEvWD_up2aT63rqSPjnps1-j9D6w,5707
+lt_tensor/model_zoo/bsc.py,sha256=OQqsQDRBf6gWqoeGeEuIaTh96AqcDyTIbO8MAMNTtI4,7045
+lt_tensor/model_zoo/disc.py,sha256=9RxyHYH2nGhxLs_yoEFVgerBfH4-qdaL2Mu9akyG0_M,5841
 lt_tensor/model_zoo/fsn.py,sha256=5ySsg2OHjvTV_coPAdZQ0f7bz4ugJB8mDYsItmd61qA,2102
 lt_tensor/model_zoo/gns.py,sha256=Tirr_grONp_FFQ_L7K-zV2lvkaC39h8mMl4QDpx9vLQ,6028
-lt_tensor/model_zoo/istft.py,sha256=0Xms2QNPAgz_ib8XTfaWl1SCHgS53oKC6-EkDkl_qe4,4863
+lt_tensor/model_zoo/istft.py,sha256=RV7KVY7q4CYzzsWXH4NGJQwSqrYWwHh-16Q62lKoA2k,3594
 lt_tensor/model_zoo/pos.py,sha256=N28v-rF8CELouYxQ9r45Jbd4ri5DNydwDgg7nzmQ4Ig,4471
 lt_tensor/model_zoo/rsd.py,sha256=5bba50g1Hm5kMexuJ4SwOIJuyQ1qJd8Acrq-Ax6CqE8,6958
 lt_tensor/model_zoo/tfrms.py,sha256=kauh-A13pk08SZ5OspEE5a-gPKD4rZr6tqMKWu3KGhk,4237
 lt_tensor/processors/__init__.py,sha256=4b9MxAJolXiJfSm20ZEspQTDm1tgLazwlPWA_jB1yLM,63
 lt_tensor/processors/audio.py,sha256=2Sta_KytTqGZh-ZeHpcCbqP6O8VT6QQVkx-7szA3Itc,8830
-lt_tensor-0.0.1a9.dist-info/licenses/LICENSE,sha256=HUnu_iSPpnDfZS_PINhO3AoVizJD1A2vee8WX7D7uXo,11358
-lt_tensor-0.0.1a9.dist-info/METADATA,sha256=4T6spgppj3ijgyjaVFt5-J6885dpOSdpJXIMnd4pQ3E,965
-lt_tensor-0.0.1a9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lt_tensor-0.0.1a9.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1a9.dist-info/RECORD,,
+lt_tensor-0.0.1a11.dist-info/licenses/LICENSE,sha256=HUnu_iSPpnDfZS_PINhO3AoVizJD1A2vee8WX7D7uXo,11358
+lt_tensor-0.0.1a11.dist-info/METADATA,sha256=DNs5JZfr_mjve_GHy13Auics3BI_f1pNYBth-dQW04M,966
+lt_tensor-0.0.1a11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lt_tensor-0.0.1a11.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1a11.dist-info/RECORD,,