lt-tensor 0.0.1a26__py3-none-any.whl → 0.0.1a27__py3-none-any.whl

lt_tensor/model_base.py

@@ -80,6 +80,62 @@ class _Devices_Base(nn.Module):
         assert isinstance(device, (str, torch.device))
         self._device = torch.device(device) if isinstance(device, str) else device
 
+    def freeze_all(self, exclude: list[str] = []):
+        for name, module in self.named_modules():
+            if name in exclude or not hasattr(module, "requires_grad"):
+                continue
+            try:
+                self.freeze_module(module)
+            except:
+                pass
+
+    def unfreeze_all(self, exclude: list[str] = []):
+        for name, module in self.named_modules():
+            if name in exclude or not hasattr(module, "requires_grad"):
+                continue
+            try:
+                self.unfreeze_module(module)
+            except:
+                pass
+
+    def freeze_module(
+        self, module_or_name: Union[str, nn.Module, nn.Parameter, "Model", Tensor]
+    ):
+        self._change_gradient_state(module_or_name, False)
+
+    def unfreeze_module(
+        self, module_or_name: Union[str, nn.Module, nn.Parameter, "Model", Tensor]
+    ):
+        self._change_gradient_state(module_or_name, True)
+
+    def _change_gradient_state(
+        self,
+        module_or_name: Union[str, nn.Module, nn.Parameter, "Model", Tensor],
+        new_state: bool,  # True = unfreeze
+    ):
+        assert isinstance(
+            module_or_name, (str, nn.Module, nn.Parameter, Model, Tensor)
+        ), f"Item '{module_or_name}' is not a valid module, parameter, tensor or a string."
+        if isinstance(module_or_name, (nn.Module, nn.Parameter, Model, Tensor)):
+            target = module_or_name
+        else:
+            target = getattr(self, module_or_name)
+
+        if isinstance(target, Tensor):
+            target.requires_grad = new_state
+        elif isinstance(target, nn.Parameter):
+            target.requires_grad = new_state
+        elif isinstance(target, Model):
+            target.freeze_all()
+        elif isinstance(target, nn.Module):
+            for param in target.parameters():
+                if hasattr(param, "requires_grad"):
+                    param.requires_grad = new_state
+        else:
+            raise ValueError(
+                f"Item '{module_or_name}' is not a valid module, parameter or tensor."
+            )
+
     def _apply_device(self):
         """Add here components that are needed to have device applied to them,
         that usually the '.to()' function fails to apply
@@ -182,20 +238,12 @@ class Model(_Devices_Base, ABC):
     """
 
     _autocast: bool = False
-    _is_unfrozen: bool = False
-    # list with modules that can be frozen or unfrozen
-    registered_freezable_modules: List[str] = []
-    is_frozen: bool = False
-    _can_be_frozen: bool = (
-        False  # to control if the module can or cannot be freezed by other modules from 'Model' class
-    )
+
     # this is to be used on the case of they module requires low-rank adapters
     _low_rank_lambda: Optional[Callable[[], nn.Module]] = (
         None  # Example: lambda: nn.Linear(32, 32, True)
     )
     low_rank_adapter: Union[nn.Identity, nn.Module, nn.Sequential] = nn.Identity()
-    # never freeze:
-    _never_freeze_modules: List[str] = ["low_rank_adapter"]
 
     # dont save list:
     _dont_save_items: List[str] = []
@@ -208,75 +256,6 @@ class Model(_Devices_Base, ABC):
     def autocast(self, value: bool):
         self._autocast = value
 
-    def freeze_all(self, exclude: Optional[List[str]] = None, force: bool = False):
-        no_exclusions = not exclude
-        no_exclusions = not exclude
-        results = []
-        for name, module in self.named_modules():
-            if (
-                name in self._never_freeze_modules
-                or not force
-                and name not in self.registered_freezable_modules
-            ):
-                results.append(
-                    (
-                        name,
-                        "Unregistered module, to freeze/unfreeze it add its name into 'registered_freezable_modules'.",
-                    )
-                )
-                continue
-            if no_exclusions:
-                self.change_frozen_state(True, module)
-            elif not any(exclusion in name for exclusion in exclude):
-                results.append((name, self.change_frozen_state(True, module)))
-            else:
-                results.append((name, "excluded"))
-        return results
-
-    def unfreeze_all(self, exclude: Optional[list[str]] = None, force: bool = False):
-        """Unfreezes all model parameters except specified layers."""
-        no_exclusions = not exclude
-        results = []
-        for name, module in self.named_modules():
-            if (
-                name in self._never_freeze_modules
-                or not force
-                and name not in self.registered_freezable_modules
-            ):
-
-                results.append(
-                    (
-                        name,
-                        "Unregistered module, to freeze/unfreeze it add it into 'registered_freezable_modules'.",
-                    )
-                )
-                continue
-            if no_exclusions:
-                self.change_frozen_state(False, module)
-            elif not any(exclusion in name for exclusion in exclude):
-                results.append((name, self.change_frozen_state(False, module)))
-            else:
-                results.append((name, "excluded"))
-        return results
-
-    def change_frozen_state(self, freeze: bool, module: nn.Module):
-        assert isinstance(module, nn.Module)
-        if module.__class__.__name__ in self._never_freeze_modules:
-            return "Not Allowed"
-        try:
-            if isinstance(module, Model):
-                if module._can_be_frozen:
-                    if freeze:
-                        return module.freeze_all()
-                    return module.unfreeze_all()
-                else:
-                    return "Not Allowed"
-            else:
-                module.requires_grad_(not freeze)
-                return not freeze
-        except Exception as e:
-            return e
-
     def 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:

lt_tensor/processors/audio.py

@@ -250,7 +250,7 @@ class AudioProcessor(Model):
             for i in range(B):
                 f0_.append(librosa.yin(self.to_numpy_safe(audio[i, :]), **yn_kwargs))
             f0 = self.from_numpy_batch(f0_, default_device, default_dtype)
-        return f0
+        return f0.squeeze()
 
     def compute_pitch_torch(
         self,
@@ -407,11 +407,7 @@ class AudioProcessor(Model):
                 mel_tensor = (
                     torch.log(eps + mel_tensor.unsqueeze(0)) - self.cfg.mean
                 ) / self.cfg.std
-            if mel_tensor.ndim == 4:
-                return mel_tensor.squeeze()
-            elif mel_tensor.ndim == 2:
-                return mel_tensor.unsqueeze(0)
-            return mel_tensor
+            return mel_tensor.squeeze()
 
         except RuntimeError as e:
             if not _recall:

{lt_tensor-0.0.1a26.dist-info → lt_tensor-0.0.1a27.dist-info}/METADATA

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

@@ -4,7 +4,7 @@ lt_tensor/losses.py,sha256=zvkCOnE5XpF3v6ymivRIdqPTsMM5zc94ZMom7YDi3zM,4946
 lt_tensor/lr_schedulers.py,sha256=LSZzqrOOLzSthD8k-W4cYPJt0vCjmHkiJkLr5e3yRTE,3659
 lt_tensor/math_ops.py,sha256=TkD4WQG42KsQ9Fg7FXOjf8f-ixtW0apf2XjaooecVx4,2257
 lt_tensor/misc_utils.py,sha256=N2r3UmxC4RM2BZBQhpjDZ_BKLrzsyIlKzopTzJbnjFU,28962
-lt_tensor/model_base.py,sha256=GvmQdt97ZSfOObBpBIq7UUTwpIE1g-aBm23za36YA0M,18431
+lt_tensor/model_base.py,sha256=DTg44N6eTXLmpIAj_ac29-M5dI_iY_sC0yA_K3E13GI,17446
 lt_tensor/monotonic_align.py,sha256=LhBd8p1xdBzg6jQrQX1j7b4PNeYGwIqM24zcU-pHOLE,2239
 lt_tensor/noise_tools.py,sha256=wFeAsHhLhSlEc5XU5LbFKaXoHeVxrWjiMeljjGdIKyM,11363
 lt_tensor/torch_commons.py,sha256=8l0bxmrAzwvyqjivCIVISXlbvKarlg4DdE0BOGSnMuQ,812
@@ -27,9 +27,9 @@ lt_tensor/model_zoo/audio_models/diffwave/__init__.py,sha256=PDuDYN1omD1RoAXcmxH
 lt_tensor/model_zoo/audio_models/hifigan/__init__.py,sha256=7GJqKLw7-juXpfp5IFzjASLut0uouDhjZ1CQknf3H68,16533
 lt_tensor/model_zoo/audio_models/istft/__init__.py,sha256=ltIuD9t1gmS3bTmCqZIwJHKrhC6DYya3OaXlskWX9kw,17606
 lt_tensor/processors/__init__.py,sha256=Pvxhh0KR65zLCgUd53_k5Z0y5JWWcO0ZBXFK9rv0o5w,109
-lt_tensor/processors/audio.py,sha256=WkumFNx8OXGQkTEU5Rkede9NLMrsGaTGY37Ti784Wv8,17028
-lt_tensor-0.0.1a26.dist-info/licenses/LICENSE,sha256=tQHc38scHOba4kDBNG4U0U6PpObaloiZG-FvKSgv2b0,11336
-lt_tensor-0.0.1a26.dist-info/METADATA,sha256=2STSK6jgD_qECwz9WygTXNDwfapEAR2mpHiS14bi9tQ,1062
-lt_tensor-0.0.1a26.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lt_tensor-0.0.1a26.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1a26.dist-info/RECORD,,
+lt_tensor/processors/audio.py,sha256=mZY7LOeYACnX8PLz5AeFe0zqEebPoN-Q44Bi3yrlZMQ,16881
+lt_tensor-0.0.1a27.dist-info/licenses/LICENSE,sha256=tQHc38scHOba4kDBNG4U0U6PpObaloiZG-FvKSgv2b0,11336
+lt_tensor-0.0.1a27.dist-info/METADATA,sha256=NpXqioPXZMvXo-HzhXrS6O1qiftDnoc8ZzOfhfUMBaY,1062
+lt_tensor-0.0.1a27.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lt_tensor-0.0.1a27.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1a27.dist-info/RECORD,,