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

lt_tensor/model_base.py

@@ -1,12 +1,14 @@
-__all__ = ["Model", "_ModelExtended", "LossTracker"]
+__all__ = ["Model", "LossTracker"]
 
 import gc
 import json
 import math
 import warnings
-from .torch_commons import *
+from lt_tensor.torch_commons import *
 from lt_utils.common import *
+from lt_tensor.misc_utils import plot_view, get_weights, get_config, updateDict
 from lt_utils.misc_utils import log_traceback, get_current_time
+from lt_utils.file_ops import load_json, save_json
 
 T = TypeVar("T")
 
@@ -29,79 +31,47 @@ class LossTracker:
             "eval": [],
         }
 
-    def append(self, loss: float, mode: Literal["train", "eval"] = "train"):
-        assert mode in self.history, f"Invalid mode '{mode}'. Use 'train' or 'eval'."
+    def append(
+        self, loss: float, mode: Union[Literal["train", "eval"], "str"] = "train"
+    ):
         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"):
+    def get(self, mode: Union[Literal["train", "eval"], "str"] = "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)
-
+        save_json(path, self.history, 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()
+            path = self.last_file
+        self.history = load_json(path, [])
+        self.last_file = path
 
+    def plot(
+        self,
+        history_keys: Union[str, List[str]],
+        max_amount: int = 0,
+        title: str = "Loss",
+    ):
+        if isinstance(history_keys, str):
+            history_keys = [history_keys]
+        return plot_view(
+            {k: v for k, v in self.history.items() if k in history_keys},
+            title,
+            max_amount,
+        )
 
-class Model(nn.Module, ABC):
-    """
-    This makes it easier to assign a device and retrieves it later
-    """
 
+class _Devices_Base(nn.Module):
     _device: torch.device = ROOT_DEVICE
     _autocast: bool = False
     _loss_history: LossTracker = LossTracker(100_000)
-    _is_unfrozen: bool = False
 
     @property
     def autocast(self):
@@ -119,7 +89,96 @@ class Model(nn.Module, ABC):
     def device(self, device: Union[torch.device, str]):
         assert isinstance(device, (str, torch.device))
         self._device = torch.device(device) if isinstance(device, str) else device
-        self._apply_device_to()
+
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(
+            *args, **kwargs
+        )
+
+        if dtype is not None:
+            if not (dtype.is_floating_point or dtype.is_complex):
+                raise TypeError(
+                    "nn.Module.to only accepts floating point or complex "
+                    f"dtypes, but got desired dtype={dtype}"
+                )
+            if dtype.is_complex:
+                warnings.warn(
+                    "Complex modules are a new feature under active development whose design may change, "
+                    "and some modules might not work as expected when using complex tensors as parameters or buffers. "
+                    "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml "
+                    "if a complex module does not work as expected."
+                )
+
+        def convert(t: Tensor):
+            try:
+                if convert_to_format is not None and t.dim() in (4, 5):
+                    return t.to(
+                        device,
+                        dtype if t.is_floating_point() or t.is_complex() else None,
+                        non_blocking,
+                        memory_format=convert_to_format,
+                    )
+                return t.to(
+                    device,
+                    dtype if t.is_floating_point() or t.is_complex() else None,
+                    non_blocking,
+                )
+            except NotImplementedError as e:
+                if str(e) == "Cannot copy out of meta tensor; no data!":
+                    raise NotImplementedError(
+                        f"{e} Please use torch.nn.Module.to_empty() instead of torch.nn.Module.to() "
+                        f"when moving module from meta to a different device."
+                    ) from None
+                else:
+                    raise
+
+        self._apply(convert)
+        self.device = device
+        return self
+
+    def _to_dvc(
+        self, device_name: str, device_id: Optional[Union[int, torch.device]] = None
+    ):
+        device = device_name
+        if device_id is not None:
+            if isinstance(device_id, Number):
+                device += ":" + str(int(device_id))
+            elif hasattr(device_id, "index"):
+                device += ":" + str(device_id.index)
+        self.device = device
+
+    def ipu(self, device: Optional[Union[int, torch.device]] = None) -> T:
+        super().ipu(device)
+        self._to_dvc("ipu", device)
+        return self
+
+    def xpu(self, device: Optional[Union[int, torch.device]] = None) -> T:
+        super().xpu(device)
+        self._to_dvc("xpu", device)
+        return self
+
+    def cuda(self, device: Optional[Union[int, torch.device]] = None) -> T:
+        super().cuda(device)
+        self._to_dvc("cuda", device)
+        return self
+
+    def mtia(self, device: Optional[Union[int, torch.device]] = None) -> T:
+        super().mtia(device)
+        self._to_dvc("mtia", device)
+        return self
+
+    def cpu(self) -> T:
+        super().cpu()
+        self._to_dvc("cpu", None)
+        return self
+
+
+class Model(_Devices_Base, ABC):
+    """
+    This makes it easier to assign a device and retrieves it later
+    """
+
+    _is_unfrozen: bool = False
 
     def _apply_device_to(self):
         """Add here components that are needed to have device applied to them,
@@ -239,103 +298,6 @@ class Model(nn.Module, ABC):
                 not_unfrozen.append((name, "excluded"))
         return dict(unfrozen=unfrozen, not_unfrozen=not_unfrozen)
 
-    def to(self, *args, **kwargs):
-        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(
-            *args, **kwargs
-        )
-
-        if dtype is not None:
-            if not (dtype.is_floating_point or dtype.is_complex):
-                raise TypeError(
-                    "nn.Module.to only accepts floating point or complex "
-                    f"dtypes, but got desired dtype={dtype}"
-                )
-            if dtype.is_complex:
-                warnings.warn(
-                    "Complex modules are a new feature under active development whose design may change, "
-                    "and some modules might not work as expected when using complex tensors as parameters or buffers. "
-                    "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml "
-                    "if a complex module does not work as expected."
-                )
-
-        def convert(t: Tensor):
-            try:
-                if convert_to_format is not None and t.dim() in (4, 5):
-                    return t.to(
-                        device,
-                        dtype if t.is_floating_point() or t.is_complex() else None,
-                        non_blocking,
-                        memory_format=convert_to_format,
-                    )
-                return t.to(
-                    device,
-                    dtype if t.is_floating_point() or t.is_complex() else None,
-                    non_blocking,
-                )
-            except NotImplementedError as e:
-                if str(e) == "Cannot copy out of meta tensor; no data!":
-                    raise NotImplementedError(
-                        f"{e} Please use torch.nn.Module.to_empty() instead of torch.nn.Module.to() "
-                        f"when moving module from meta to a different device."
-                    ) from None
-                else:
-                    raise
-
-        self._apply(convert)
-        self.device = device
-        self._apply_device_to()
-        return self
-
-    def ipu(self, device: Optional[Union[int, torch.device]] = None) -> T:
-        super().ipu(device)
-        dvc = "ipu"
-        if device is not None:
-            dvc += (
-                ":" + 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:
-        super().xpu(device)
-        dvc = "xpu"
-        if device is not None:
-            dvc += (
-                ":" + 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:
-        super().cuda(device)
-        dvc = "cuda"
-        if device is not None:
-            dvc += (
-                ":" + 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:
-        super().mtia(device)
-        dvc = "mtia"
-        if device is not None:
-            dvc += (
-                ":" + 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):
         """Gets the number of trainable parameters from either the entire model or from a specific module."""
         if module_name is not None:
@@ -411,6 +373,19 @@ class Model(nn.Module, ABC):
         else:
             print(f"Non-Trainable Parameters: {params}")
 
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: Union[str, PathLike],
+        model_config: Optional[Union[str, PathLike]] = None,
+        model_config_dict: Optional[Dict[str, Any]] = None,
+    ):
+        """TODO: Finish this implementation or leave it as an optional function."""
+        assert isinstance(model_path, (str, PathLike, bytes))
+        model_path = get_weights(model_path)
+        model_config = get_config(model_config, model_config_dict)
+        return model_path, model_config
+
     def save_weights(
         self,
         path: Union[Path, str],
@@ -494,7 +469,9 @@ class Model(nn.Module, ABC):
         pass
 
     def add_loss(
-        self, loss: Union[float, list[float]], mode: Literal["train", "eval"] = "train"
+        self,
+        loss: Union[float, list[float]],
+        mode: Union[Literal["train", "eval"]] = "train",
     ):
         if isinstance(loss, Number) and loss:
             self._loss_history.append(loss, mode)
@@ -513,8 +490,12 @@ class Model(nn.Module, ABC):
     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")
+    def get_loss_avg(
+        self,
+        mode: Union[Literal["train", "eval"], str],
+        quantity: int = 0,
+    ):
+        t_list = self._loss_history.get(mode)
         if not t_list:
             return float("nan")
         if quantity > 0:
@@ -524,9 +505,10 @@ class Model(nn.Module, ABC):
     def freeze_unfreeze_loss(
         self,
         losses: Optional[Union[float, List[float]]] = None,
-        trigger_loss: float = 0.1,
+        trigger_loss: Union[float, bool] = 0.1,
         excluded_modules: Optional[List[str]] = None,
-        eval_last: int = 1000,
+        max_items: int = 1000,
+        loss_name: str = "train",
     ):
         """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
@@ -534,18 +516,29 @@ class Model(nn.Module, ABC):
 
         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.
+            trigger_loss (float, bool, optional): The value where the weights will be either freeze or unfreeze. If set to a boolean it will freeze or unfreeze immediately according to the value (True = Freeze, False = 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.
-
+            max_items (float, optional): The number of previous losses to be locked behind to calculate the current average. Default to 1000.
+            loss_name (str, optional): Responsible to define with key to recover the loss.
         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 isinstance(trigger_loss, bool):
+            if 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
+
+        value = self.get_loss_avg(loss_name, max_items)
 
         if value <= trigger_loss:
             if self._is_unfrozen:
@@ -557,42 +550,3 @@ class Model(nn.Module, ABC):
                 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/__init__.py

@@ -1,11 +1,20 @@
 __all__ = [
-    "bsc",  # basic
-    "rsd",  # residual
-    "tfrms",  # transformer
-    "pos",  # positional encoders
-    "fsn",  # fusion
-    "gns",  # generators
-    "disc", # discriminators
-    "istft" # self-explanatory
+    "basic",  # basic
+    "residual",  # residual
+    "transformer",  # transformer
+    "pos_encoder", 
+    "fusion", 
+    "features", 
+    "discriminator",  
+    "istft",
 ]
-from . import bsc, fsn, gns, istft, pos, rsd, tfrms, disc
+from . import (
+    basic,
+    discriminator,
+    features,
+    fusion,
+    istft,
+    pos_encoder,
+    residual,
+    transformer,
+)

lt_tensor/model_zoo/{bsc.py → basic.py}

@@ -8,11 +8,14 @@ __all__ = [
     "StyleEncoder",
     "PatchEmbed1D",
     "MultiScaleEncoder1D",
+    "UpSampleConv1D",
 ]
-
-from ..torch_commons import *
-from ..model_base import Model
-from ..transform import get_sinusoidal_embedding
+import torch.nn.functional as F
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
+from lt_tensor.transform import get_sinusoidal_embedding
+from lt_utils.common import *
+import math
 
 
 class FeedForward(Model):
@@ -30,6 +33,9 @@ class FeedForward(Model):
             nn.Linear(d_model, ff_dim),
             activation,
             nn.Dropout(dropout),
+            nn.Linear(ff_dim, ff_dim),
+            activation,
+            nn.Dropout(dropout),
             nn.Linear(ff_dim, d_model),
             normalizer,
         )
@@ -38,6 +44,24 @@ class FeedForward(Model):
         return self.net(x)
 
 
+class UpSampleConv1D(nn.Module):
+    def __init__(self, upsample: bool = False, dim_in: int = 0, dim_out: int = 0):
+        super().__init__()
+        if upsample:
+            self.upsample = lambda x: F.interpolate(x, scale_factor=2, mode="nearest")
+        else:
+            self.upsample = nn.Identity()
+
+        if dim_in == dim_out:
+            self.learned = nn.Identity()
+        else:
+            self.learned = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def forward(self, x):
+        x = self.upsample(x)
+        return self.learned(x)
+
+
 class MLP(Model):
     def __init__(
         self,
@@ -161,7 +185,6 @@ class StyleEncoder(Model):
         self.linear = nn.Linear(hidden, out_dim)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x: [B, Mels, T]
         x = self.net(x).squeeze(-1)  # [B, hidden]
         return self.linear(x)  # [B, out_dim]
 
@@ -208,10 +231,10 @@ 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):
+    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),
@@ -230,3 +253,96 @@ class AudioClassifier(Model):
 
     def forward(self, x):
         return self.model(x)
+
+
+class LoRALinearLayer(nn.Module):
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        rank: int = 4,
+        alpha: float = 1.0,
+    ):
+        super().__init__()
+        self.down = nn.Linear(in_features, rank, bias=False)
+        self.up = nn.Linear(rank, out_features, bias=False)
+        self.alpha = alpha
+        self.rank = rank
+        self.out_features = out_features
+        self.in_features = in_features
+        self.ah = self.alpha / self.rank
+        self._down_dt = self.down.weight.dtype
+
+        nn.init.normal_(self.down.weight, std=1 / rank)
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        orig_dtype = hidden_states.dtype
+        down_hidden_states = self.down(hidden_states.to(self._down_dt))
+        up_hidden_states = self.up(down_hidden_states) * self.ah
+        return up_hidden_states.to(orig_dtype)
+
+
+class LoRAConv1DLayer(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        kernel_size: Union[int, Tuple[int, ...]] = 1,
+        rank: int = 4,
+        alpha: float = 1.0,
+    ):
+        super().__init__()
+        self.down = nn.Conv1d(
+            in_features, rank, kernel_size, padding=kernel_size // 2, bias=False
+        )
+        self.up = nn.Conv1d(
+            rank, out_features, kernel_size, padding=kernel_size // 2, bias=False
+        )
+        self.ah = alpha / rank
+        self._down_dt = self.down.weight.dtype
+        nn.init.kaiming_uniform_(self.down.weight, a=math.sqrt(5))
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        orig_dtype = inputs.dtype
+        down_hidden_states = self.down(inputs.to(self._down_dt))
+        up_hidden_states = self.up(down_hidden_states) * self.ah
+        return up_hidden_states.to(orig_dtype)
+
+
+class LoRAConv2DLayer(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        kernel_size: Union[int, Tuple[int, ...]] = (1, 1),
+        rank: int = 4,
+        alpha: float = 1.0,
+    ):
+        super().__init__()
+        self.down = nn.Conv2d(
+            in_features,
+            rank,
+            kernel_size,
+            padding="same",
+            bias=False,
+        )
+        self.up = nn.Conv2d(
+            rank,
+            out_features,
+            kernel_size,
+            padding="same",
+            bias=False,
+        )
+        self.ah = alpha / rank
+
+        nn.init.kaiming_normal_(self.down.weight, a=0.2)
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        orig_dtype = inputs.dtype
+        down_hidden_states = self.down(inputs.to(self._down_dt))
+        up_hidden_states = self.up(down_hidden_states) * self.ah
+        return up_hidden_states.to(orig_dtype)

lt_tensor/model_zoo/{disc.py → discriminator.py}

@@ -1,4 +1,4 @@
-from ..torch_commons import *
+from lt_tensor.torch_commons import *
 import torch.nn.functional as F
 from lt_tensor.model_base import Model
 from lt_utils.common import *