lt-tensor 0.0.1.dev0__py3-none-any.whl

lt_tensor/_basics.py

@@ -0,0 +1,244 @@
+__all__ = ["Model"]
+
+
+import warnings
+from ._torch_commons import *
+
+ROOT_DEVICE = torch.device(torch.zeros(1).device)
+
+class _ModelDevice(nn.Module):
+    """
+    This makes it easier to assign a device and retrieves it later
+    """
+
+    _device: torch.device = ROOT_DEVICE
+
+    @property
+    def device(self):
+        return self._device
+
+    @device.setter
+    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.tp_apply_device_to()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def tp_apply_device_to(self):
+        """Add here components that are needed to have device applied to them,
+        that usualy the '.to()' function fails to apply
+
+        example:
+        ```
+        def tp_apply_device_to(self):
+            self.my_tensor = self.my_tensor.to(device=self.device)
+        ```
+        """
+        pass
+
+    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 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
+        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
+        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
+        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
+        return self
+
+    def cpu(self) -> T:
+        super().cpu()
+        self.device = "cpu"
+        return self
+
+
+class Model(_ModelDevice, ABC):
+    """
+    For easier access of items in a model
+    This applies to all except to "forward", "save_weights", "load_weights",
+    "train_step", and "inference".
+    """
+
+    _autocast: bool = False
+
+    @property
+    def autocast(self):
+        return self._autocast
+
+    @autocast.setter
+    def autocast(self, value: bool):
+        self._autocast = value
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        device = kwargs.get("device", None)
+        if isinstance(device, (torch.device, str)):
+            self.to(device)
+        self.autocast = kwargs.get("autocast", self.autocast)
+
+    def tp_count_parameters(self, module_name: Optional[str] = None):
+        """Gets the numer of parameters from either the entire model or from a specific module"""
+        if module_name is not None:
+            assert hasattr(self, module_name), f"Module {module_name} does not exits"
+            module = getattr(self, module_name)
+            return sum([x.numel() for x in module.parameters()])
+        else:
+            return sum([x.numel() for x in self.parameters()])
+
+    def tp_get_weights(self, module_name: Optional[str] = None) -> List[Tensor]:
+        """Returns the weights of the model entrie model or from a specified module"""
+        if module_name is not None:
+            assert hasattr(self, module_name), f"Module {module_name} does not exits"
+            module = getattr(self, module_name)
+            return [x.data.detach() for x in module.parameters()]
+        return [x.data.detach() for x in self.parameters()]
+
+    def tp_print_num_params(self, module_name: Optional[str] = None) -> List[Tensor]:
+        params = format(self.tp_count_parameters(module_name), ",").replace(",", ".")
+
+        if module_name:
+            print(f'Parameters from "{module_name}": {params}')
+        else:
+            print(f"Parameters: {params}")
+
+    def save_weights(self, path: Union[Path, str]):
+        path = Path(path)
+        if path.exists():
+            assert (
+                path.is_file()
+            ), "The provided path exists but its a directory not a file!"
+            path.rmdir()
+        path.parent.mkdir(exist_ok=True, parents=True)
+        torch.save(obj=self.state_dict(), f=str(path))
+
+    def load_weights(
+        self,
+        path: Union[Path, str],
+        raise_if_not_exists: bool = False,
+        strict: bool = True,
+        assign: bool = False,
+        weights_only: bool = False,
+        mmap: Optional[bool] = None,
+        **torch_loader_kwargs,
+    ):
+        path = Path(path)
+        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!"
+        state_dict = torch.load(
+            str(path), weights_only=weights_only, mmap=mmap, **torch_loader_kwargs
+        )
+        incompatible_keys = self.load_state_dict(
+            state_dict,
+            strict=strict,
+            assign=assign,
+        )
+        return incompatible_keys
+
+    @torch.no_grad()
+    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,
+        **kwargs,
+    ):
+        """Train Step"""
+        if not self.training:
+            self.train()
+        return self(*args, **kwargs)
+
+    @abstractmethod
+    def forward(
+        self, *args, **kwargs
+    ) -> Union[Tensor, Sequence[Tensor], Dict[Any, Union[Any, Tensor]]]:
+        pass

lt_tensor/_torch_commons.py

@@ -0,0 +1,12 @@
+import torch
+from torch import nn, optim
+import torch.nn.functional as F
+from torch.optim import Optimizer
+from torch.nn import Module, L1Loss, MSELoss
+from torch.nn.utils import remove_weight_norm
+from torch import Tensor, FloatTensor, device, LongTensor
+from torch.nn.utils.parametrizations import weight_norm, spectral_norm
+
+from lt_utils.common import *
+
+DeviceType: TypeAlias = Union[device, str]

lt_tensor/lr_schedulers.py

@@ -0,0 +1,108 @@
+import math
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+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 WaveringLRScheduler(_LRScheduler):
+    def __init__(
+        self, optimizer, base_lr, max_lr, period=1000, decay=0.999, last_epoch=-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
+        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]

lt_tensor/math_ops.py

@@ -0,0 +1,120 @@
+from ._torch_commons import *
+
+
+def sin_tensor(x: Tensor, freq: float = 1.0) -> Tensor:
+    """Applies sine function element-wise."""
+    return torch.sin(x * freq)
+
+
+def cos_tensor(x: Tensor, freq: float = 1.0) -> Tensor:
+    """Applies cosine function element-wise."""
+    return torch.cos(x * freq)
+
+
+def sin_plus_cos(x: Tensor, freq: float = 1.0) -> Tensor:
+    """Returns sin(x) + cos(x)."""
+    return torch.sin(x * freq) + torch.cos(x * freq)
+
+
+def sin_times_cos(x: Tensor, freq: float = 1.0) -> Tensor:
+    """Returns sin(x) * cos(x)."""
+    return torch.sin(x * freq) * torch.cos(x * freq)
+
+
+def apply_window(x: Tensor, window_type: str = "hann") -> Tensor:
+    """Applies a window function to a 1D tensor."""
+    if window_type == "hann":
+        window = torch.hann_window(x.shape[-1], device=x.device)
+    elif window_type == "hamming":
+        window = torch.hamming_window(x.shape[-1], device=x.device)
+    else:
+        raise ValueError(f"Unsupported window type: {window_type}")
+    return x * window
+
+
+def shift_ring(x: Tensor, dim: int = -1) -> Tensor:
+    """Circularly shifts tensor values: last becomes first (along given dim)."""
+    return torch.roll(x, shifts=1, dims=dim)
+
+
+def dot_product(x: Tensor, y: Tensor, dim: int = -1) -> Tensor:
+    """Computes dot product along the specified dimension."""
+    return torch.sum(x * y, dim=dim)
+
+
+def normalize_tensor(x: Tensor, eps: float = 1e-8) -> Tensor:
+    """Normalizes a tensor to unit norm (L2)."""
+    return x / (torch.norm(x, dim=-1, keepdim=True) + eps)
+
+
+def stft(
+    waveform: Tensor,
+    n_fft: int = 512,
+    hop_length: Optional[int] = None,
+    win_length: Optional[int] = None,
+    window_fn: str = "hann",
+    center: bool = True,
+    return_complex: bool = True,
+) -> Tensor:
+    """Performs short-time Fourier transform using PyTorch."""
+    window = (
+        torch.hann_window(win_length or n_fft).to(waveform.device)
+        if window_fn == "hann"
+        else None
+    )
+    return torch.stft(
+        input=waveform,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=window,
+        center=center,
+        return_complex=return_complex,
+    )
+
+
+def istft(
+    stft_matrix: Tensor,
+    n_fft: int = 512,
+    hop_length: Optional[int] = None,
+    win_length: Optional[int] = None,
+    window_fn: str = "hann",
+    center: bool = True,
+    length: Optional[int] = None,
+) -> Tensor:
+    """Performs inverse short-time Fourier transform using PyTorch."""
+    window = (
+        torch.hann_window(win_length or n_fft).to(stft_matrix.device)
+        if window_fn == "hann"
+        else None
+    )
+    return torch.istft(
+        input=stft_matrix,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=window,
+        center=center,
+        length=length,
+    )
+
+
+def fft(x: Tensor, norm: Optional[str] = "backward") -> Tensor:
+    """Returns the FFT of a real tensor."""
+    return torch.fft.fft(x, norm=norm)
+
+
+def ifft(x: Tensor, norm: Optional[str] = "backward") -> Tensor:
+    """Returns the inverse FFT of a complex tensor."""
+    return torch.fft.ifft(x, norm=norm)
+
+
+def log_magnitude(stft_complex: Tensor, eps: float = 1e-5) -> Tensor:
+    """Returns log magnitude from complex STFT."""
+    magnitude = torch.abs(stft_complex)
+    return torch.log(magnitude + eps)
+
+
+def phase(stft_complex: Tensor) -> Tensor:
+    """Returns phase from complex STFT."""
+    return torch.angle(stft_complex)