lucid-dl 2.7.5__py3-none-any.whl → 2.7.7__py3-none-any.whl

lucid/nn/__init__.py

@@ -4,3 +4,4 @@ from lucid.nn.modules import *
 from lucid.nn.fused import *
 
 import lucid.nn.init as init
+import lucid.nn.util as util

lucid/nn/module.py

@@ -118,7 +118,7 @@ class Module:
         for param in self.parameters():
             param.zero()
 
-    def forward(self, *args: Any, **kwargs: Any) -> Tensor | tuple[Tensor, ...]:
+    def forward(self) -> Tensor | tuple[Tensor, ...]:
         raise NotImplementedError(
             "The forward method must be implemented by the subclass."
         )
@@ -204,7 +204,7 @@ class Module:
                 destination=destination, prefix=prefix + name + ".", keep_vars=keep_vars
             )
 
-        for key in destination.keys():
+        for key in list(destination.keys()):
             if key in self._state_dict_pass_attr:
                 del destination[key]
 
@@ -229,10 +229,8 @@ class Module:
             if key in own_state:
                 attr = own_state[key]
                 if isinstance(attr, (nn.Parameter, nn.Buffer)):
-                    if isinstance(value, Tensor):
-                        attr.data = value.data
-                    else:
-                        attr.data = value
+                    value_t = Tensor(value, device=self.device)
+                    attr.data = value_t.data
                 else:
                     setattr(self, key, value)
             elif strict:

lucid/nn/util.py

@@ -0,0 +1,60 @@
+from typing import Iterable
+
+import lucid
+
+from lucid._tensor import Tensor
+from lucid.types import _Scalar
+
+
+__all__ = ["grad_norm", "clip_grad_norm", "clip_grad_value"]
+
+
+def _as_iter(parameters: Iterable[Tensor] | Tensor) -> list[Tensor]:
+    if isinstance(parameters, Tensor):
+        return [parameters]
+    return list(parameters)
+
+
+def grad_norm(parameters: Iterable[Tensor] | Tensor, norm_type: int = 2) -> Tensor:
+    parameters = _as_iter(parameters)
+    device = parameters[0].device
+
+    params: list[Tensor] = [p for p in parameters if p.grad is not None]
+    if not params:
+        return Tensor(0.0, device=device)
+
+    norm_pow_sum = 0.0
+    for p in params:
+        param_norm = lucid.linalg.norm(lucid.ravel(p.grad), ord=norm_type).item()
+        norm_pow_sum += param_norm**norm_type
+
+    total_norm = norm_pow_sum ** (1.0 / norm_type)
+    return Tensor(total_norm, device=device)
+
+
+def clip_grad_norm(
+    parameters: Iterable[Tensor] | Tensor,
+    max_norm: _Scalar,
+    norm_type: int = 2,
+    eps: float = 1e-7,
+) -> float:
+    params: list[Tensor] = [p for p in _as_iter(parameters) if p.grad is not None]
+    total_norm = grad_norm(params, norm_type=norm_type)
+
+    clip_coef = float(max_norm) / (total_norm.item() + eps)
+    if clip_coef < 1.0:
+        for p in params:
+            p.grad = p.grad * clip_coef
+
+    return total_norm
+
+
+def clip_grad_value(parameters: Iterable[Tensor] | Tensor, clip_value: _Scalar) -> None:
+    params = [p for p in _as_iter(parameters) if p.grad is not None]
+    if not params:
+        return
+
+    lo, hi = -float(clip_value), float(clip_value)
+    for p in params:
+        g_clip = lucid.clip(p.grad, lo, hi).data
+        p.grad = g_clip

lucid/optim/_base.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Any, Iterable, OrderedDict
+from typing import Any, Iterable
 from abc import ABC, abstractmethod
 import copy
 
@@ -12,19 +12,19 @@ class Optimizer(ABC):
     def __init__(
         self, params: Iterable[nn.Parameter], defaults: dict[str, Any]
     ) -> None:
-        super().__init__()
         if not isinstance(params, Iterable):
             raise TypeError("params should be an iterable of Parameters.")
 
-        params = list(params)
-        self.param_groups = self.param_groups_setup(params, defaults)
-        self.defaults = defaults
-        self.state: dict[nn.Parameter, dict[str, Any]] = defaultdict(dict)
+        param_list = list(params)
+        for p in param_list:
+            if not isinstance(p, nn.Parameter):
+                raise TypeError(f"Expected nn.Parameter, got {type(p).__name__}.")
 
-    def param_groups_setup(
-        self, params: list[nn.Parameter], defaults: dict[str, Any]
-    ) -> list[dict[str, Any]]:
-        return [{"params": list(params), **defaults}]
+        self.defaults: dict[str, Any] = dict(defaults)
+        self.param_groups: list[dict[str, Any]] = self.param_groups_setup(
+            param_list, self.defaults
+        )
+        self.state: dict[nn.Parameter, dict[str, Any]] = defaultdict(dict)
 
     @abstractmethod
     def step(self, closure: _OptimClosure | None = None) -> Any | None:
@@ -33,25 +33,110 @@ class Optimizer(ABC):
     def zero_grad(self) -> None:
         for group in self.param_groups:
             for param in group["params"]:
-                param.zero_grad()
+                if isinstance(param, nn.Parameter):
+                    param.zero_grad()
+
+    def param_groups_setup(
+        self, params: list[nn.Parameter], defaults: dict[str, Any]
+    ) -> list[dict[str, Any]]:
+        return [{"params": list(params), **defaults}]
 
     def add_param_group(self, param_group: dict[str, Any]) -> None:
-        for group in self.param_groups:
-            if set(group["params"]) & set(param_group["params"]):
-                raise ValueError(
-                    "Some parameters appear in more than one parameter group."
+        if "params" not in param_group:
+            raise ValueError("param_group must have a 'params' key.")
+
+        params = list(param_group["params"])
+        if len(params) == 0:
+            raise ValueError("param_group['params'] must be non-empty.")
+
+        for p in params:
+            if not isinstance(p, nn.Parameter):
+                raise TypeError(
+                    f"Expected nn.Parameter in param_group, got {type(p).__name__}."
                 )
-        self.param_groups.append(param_group)
 
-    def state_dict(self) -> OrderedDict:
-        return {
-            "state": copy.deepcopy(self.state),
-            "param_groups": copy.deepcopy(self.param_groups),
+        existing = set()
+        for g in self.param_groups:
+            existing.update(g["params"])
+
+        if any(p in existing for p in params):
+            raise ValueError("Some parameters appear in more than one parameter group.")
+
+        filled = {
+            **self.defaults,
+            **{k: v for k, v in param_group.items() if k != "params"},
         }
+        filled["params"] = params
+        self.param_groups.append(filled)
+
+    def _flat_params(self) -> list[nn.Parameter]:
+        flat: list[nn.Parameter] = []
+        for g in self.param_groups:
+            flat.extend(g["params"])
+
+        return flat
+
+    def state_dict(self) -> dict:
+        param_to_idx: dict[nn.Parameter, int] = {}
+        for idx, p in enumerate(self._flat_params()):
+            if p not in param_to_idx:
+                param_to_idx[p] = idx
+
+        packed_state: dict[int, dict[str, Any]] = {}
+        for p, st in self.state.items():
+            if p in param_to_idx:
+                packed_state[param_to_idx[p]] = copy.deepcopy(st)
+
+        packed_groups: list[dict[str, Any]] = []
+        for g in self.param_groups:
+            new_g: dict[str, Any] = {}
+
+            for k, v in g.items():
+                if k == "params":
+                    new_g[k] = [param_to_idx[p] for p in v]
+                else:
+                    new_g[k] = copy.deepcopy(v)
+
+            packed_groups.append(new_g)
+
+        return {"state": packed_state, "param_groups": packed_groups}
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        if (
+            not isinstance(state_dict, dict)
+            or "state" not in state_dict
+            or "param_groups" not in state_dict
+        ):
+            raise TypeError("Invalid state_dict format for Optimizer.")
+
+        saved_groups = state_dict["param_groups"]
+        saved_state = state_dict["state"]
+
+        current_params = self._flat_params()
+        n_current = len(current_params)
+
+        new_groups: list[dict[str, Any]] = []
+        for sg in saved_groups:
+            if "params" not in sg:
+                raise KeyError("Saved param_group missing 'params'.")
+            indices: list[int] = list(sg["params"])
+
+            if any(i < 0 or i >= n_current for i in indices):
+                raise IndexError("Saved state refers to parameter index out of range.")
+
+            params = [current_params[i] for i in indices]
+            ng = {
+                k: (params if k == "params" else copy.deepcopy(v))
+                for k, v in sg.items()
+            }
+            new_groups.append(ng)
+
+        self.param_groups = new_groups
 
-    def load_state_dict(self, state_dict: OrderedDict) -> None:
-        self.state = defaultdict(dict, copy.deepcopy(state_dict["state"]))
-        self.param_groups = copy.deepcopy(state_dict["param_groups"])
+        self.state = defaultdict(dict)
+        for i, p in enumerate(self._flat_params()):
+            if i in saved_state:
+                self.state[p] = copy.deepcopy(saved_state[i])
 
     def __repr__(self) -> str:
-        return f"{type(self).__name__}({self.param_groups})"
+        return f"{type(self).__name__}({self.defaults})"

lucid/optim/lr_scheduler/_base.py

@@ -8,56 +8,78 @@ class LRScheduler(ABC):
     def __init__(
         self, optimizer: Optimizer, last_epoch: int = -1, verbose: bool = False
     ) -> None:
-        super().__init__()
         if not hasattr(optimizer, "param_groups"):
             raise TypeError(f"{type(optimizer).__name__} is not a valid optimizer.")
 
         self.optimizer = optimizer
         self.last_epoch = last_epoch
         self.verbose = verbose
-        self.base_lrs = [group["lr"] for group in optimizer.param_groups]
+        self.base_lrs: list[float] = [float(g["lr"]) for g in optimizer.param_groups]
 
         self._step_count = 0
-        self._last_lr = [group["lr"] for group in optimizer.param_groups]
+        self._last_lr: list[float] = [float(g["lr"]) for g in optimizer.param_groups]
 
     @abstractmethod
     def get_lr(self) -> list[float]:
-        raise NotImplementedError("get_lr must be implemented in subclasses.")
+        raise NotImplementedError
 
     def step(self, epoch: int | None = None) -> None:
-        if epoch is not None:
-            self.last_epoch = epoch
+        if epoch is None:
+            self.last_epoch += 1
         else:
-            self._step_count += 1
-            self.last_epoch = self._step_count
+            self.last_epoch = int(epoch)
+        self._step_count += 1
 
         new_lrs = self.get_lr()
-        for param_group, lr in zip(self.optimizer.param_groups, new_lrs):
-            param_group["lr"] = lr
+        if len(new_lrs) != len(self.optimizer.param_groups):
+            raise ValueError(
+                f"get_lr returned {len(new_lrs)} values, "
+                f"but optimizer has {len(self.optimizer.param_groups)} param groups."
+            )
 
-        self._last_lr = new_lrs
+        for group, lr in zip(self.optimizer.param_groups, new_lrs):
+            group["lr"] = float(lr)
+
+        self._last_lr = [float(g["lr"]) for g in self.optimizer.param_groups]
 
         if self.verbose:
-            print(f"Epoch {self.last_epoch}: setting learning rates to {new_lrs}.")
+            print(
+                f"Epoch {self.last_epoch}: setting learning rates to {self._last_lr}."
+            )
 
     def state_dict(self) -> dict[str, Any]:
         return {
-            "last_epoch": self.last_epoch,
-            "base_lrs": self.base_lrs,
-            "_step_count": self._step_count,
-            "_last_lr": self._last_lr,
+            "last_epoch": int(self.last_epoch),
+            "base_lrs": [float(x) for x in self.base_lrs],
+            "_step_count": int(self._step_count),
+            "_last_lr": [float(x) for x in self._last_lr],
+            "_group_count": len(self.optimizer.param_groups),
         }
 
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
-        self.last_epoch = state_dict["last_epoch"]
-        self.base_lrs = state_dict["base_lrs"]
+        required = {"last_epoch", "base_lrs", "_step_count", "_last_lr"}
+        missing = required - set(state_dict)
+        if missing:
+            raise KeyError(f"Missing keys in scheduler state_dict: {missing}")
+
+        saved_group_count = int(
+            state_dict.get("_group_count", len(state_dict["_last_lr"]))
+        )
+        current_group_count = len(self.optimizer.param_groups)
+        if saved_group_count != current_group_count:
+            raise ValueError(
+                "Cannot load scheduler state: param group count mismatch "
+                f"(saved={saved_group_count}, current={current_group_count})."
+            )
 
-        self._step_count = state_dict["_step_count"]
-        self._last_lr = state_dict["_last_lr"]
+        self.last_epoch = int(state_dict["last_epoch"])
+        self.base_lrs = [float(x) for x in state_dict["base_lrs"]]
+        self._step_count = int(state_dict["_step_count"])
+        self._last_lr = [float(x) for x in state_dict["_last_lr"]]
 
-        for param_group, lr in zip(self.optimizer.param_groups, self._last_lr):
-            param_group["lr"] = lr
+        for group, lr in zip(self.optimizer.param_groups, self._last_lr):
+            group["lr"] = float(lr)
 
     @property
     def last_lr(self) -> list[float]:
-        return self._last_lr
+        return list(self._last_lr)

lucid/port.py

@@ -5,14 +5,13 @@ from typing import Literal
 
 from lucid._tensor import Tensor
 from lucid.nn import Module
-from lucid.types import _NumPyArray
 
 
 __all__ = ["save", "load"]
 
-_LucidPortable = Tensor | Module | OrderedDict
+_LucidPortable = Tensor | Module | OrderedDict | dict
 
-FORMAT_VERSION: float = 1.0
+FORMAT_VERSION: float = 1.1
 
 EXTENSIONS = Literal[".lct", ".lcd", ".safetensors"]
 
@@ -30,7 +29,7 @@ def save(obj: _LucidPortable, path: Path | str, safetensors: bool = False) -> Pa
     if path.suffix == "":
         if isinstance(obj, Tensor):
             path = path.with_suffix(".lct")
-        elif isinstance(obj, (Module, OrderedDict)):
+        elif isinstance(obj, (Module, OrderedDict, dict)):
             path = (
                 path.with_suffix(".safetensors")
                 if safetensors
@@ -56,10 +55,14 @@ def save(obj: _LucidPortable, path: Path | str, safetensors: bool = False) -> Pa
     elif suffix == ".lcd":
         if isinstance(obj, Module):
             obj = obj.state_dict()
-        if not isinstance(obj, OrderedDict):
-            raise TypeError("Expected a state_dict (OrderedDict) for .lcd file.")
+        if not isinstance(obj, (OrderedDict, dict)):
+            raise TypeError("Expected a state_dict for .lcd file.")
 
-        data = {"type": "OrderedDict", "format_version": FORMAT_VERSION, "content": obj}
+        data = {
+            "type": type(obj).__name__,
+            "format_version": FORMAT_VERSION,
+            "content": obj,
+        }
 
     elif suffix == ".safetensors":
         try:
@@ -72,10 +75,8 @@ def save(obj: _LucidPortable, path: Path | str, safetensors: bool = False) -> Pa
 
         if isinstance(obj, Module):
             obj = obj.state_dict()
-        if not isinstance(obj, OrderedDict):
-            raise TypeError(
-                "Expected a state_dict (OrderedDict) for .safetensors file."
-            )
+        if not isinstance(obj, (OrderedDict, dict)):
+            raise TypeError("Expected a state_dict for .safetensors file.")
 
         save_file(obj, str(path))
         return path.resolve()
@@ -122,7 +123,7 @@ def load(path: Path | str) -> _LucidPortable:
         array = data["content"]
         return Tensor(array)
 
-    elif file_type == "OrderedDict":
+    elif file_type in {"OrderedDict", "dict"}:
         return data["content"]
 
     else:

{lucid_dl-2.7.5.dist-info → lucid_dl-2.7.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lucid-dl
-Version: 2.7.5
+Version: 2.7.7
 Summary: Lumerico's Comprehensive Interface for Deep Learning
 Home-page: https://github.com/ChanLumerico/lucid
 Author: ChanLumerico

{lucid_dl-2.7.5.dist-info → lucid_dl-2.7.7.dist-info}/RECORD

@@ -1,6 +1,6 @@
 lucid/__init__.py,sha256=NzXUnsIgDEq0Yx146zxhqID44X_0yNirkxt68NuL6Kg,7900
 lucid/error.py,sha256=qnTiVuZm3c5-DIt-OOyobZ7RUm7E1K4NR0j998LG1ug,709
-lucid/port.py,sha256=Ld5h6F3U6qNA-PSpxuPhbYeyqhZhMVCfTBXaT0xLcW0,3814
+lucid/port.py,sha256=Kt1YaSWef_eKF4KRj-UFhirvFC5urEESfYQ_BSlBZGE,3811
 lucid/types.py,sha256=3yj99eZNv8N3bupP4htHU_SEvrqNyRuq2k7WGLVD-X0,3605
 lucid/_backend/__init__.py,sha256=n1bnYdeb_bNDBKASWGywTRa0Ne9hMAkal3AuVZJgovI,5
 lucid/_backend/core.py,sha256=xMyhNoEswIxCr69wTNo3QXtNVUfX7xsmDP4t0JRlRzQ,5798
@@ -71,10 +71,11 @@ lucid/models/objdet/yolo/yolo_v3.py,sha256=B5U42Npwfg8nSgU9E261zf0cbQS9RVYrX1ADD
 lucid/models/objdet/yolo/yolo_v4.py,sha256=RFbBumreXmy6s8IYZvUuhW0893ss8sx_8Vgi6KbBKWo,21467
 lucid/models/seq2seq/__init__.py,sha256=wjsrhj4H_AcqwwbebAN8b68QBA8L6p1_12dkG2995-w,27
 lucid/models/seq2seq/transformer.py,sha256=y5rerCs1s6jXTsVvbgscWScKpQKuSu1fezsBe7PNTRA,3513
-lucid/nn/__init__.py,sha256=Kc6_wlpWo0_AtywX8aEWtzjKb0ju2c2cKGNsEY9ho4E,153
+lucid/nn/__init__.py,sha256=_hk6KltQIJuWXowXstMSu3TjiaTP8zMLNvGpjnA9Mpw,182
 lucid/nn/fused.py,sha256=ZGOQmDThaGNQLC59y3M7s993K_K09ce6IZP8cFX8FUE,5498
-lucid/nn/module.py,sha256=PeMVEdGYg5bVgRVQUdoTTT5zrZ-irt2vyuWQN6xCaBI,22048
+lucid/nn/module.py,sha256=XvFWJ8NqXeZpr3RmKBQBz5eqT535Oi_7DaPN1Zi9gJc,21971
 lucid/nn/parameter.py,sha256=jDaWukWecCcH9ri65SefNls66MmyTyucFolWbzSjapc,856
+lucid/nn/util.py,sha256=Yw1iBSPrGV_r_F51qpqLYdafNE_hyaA0DPWYP-rjaig,1699
 lucid/nn/functional/__init__.py,sha256=90Zi7jClPOiiSYx-Qkg0QTideKD6GigbWON9eFCoxzg,13869
 lucid/nn/functional/_activation.py,sha256=nQVwArvPuwkUpLMLCNABTw96Zgw9VsPB8SyXCL6t2LM,1331
 lucid/nn/functional/_attention.py,sha256=nrZF3-2AR03kNo1PGNszujhWlAVcab_FNQwOCWZT47I,946
@@ -102,13 +103,13 @@ lucid/nn/modules/sparse.py,sha256=EpjiviED2nI55wUjh1twFwa4Lvlrzw0TR6lpCDGeSbo,11
 lucid/nn/modules/transformer.py,sha256=z56emF_eX18pxRELjfmmsY-7Bn9h2yjIdxCaxs6YDwA,11246
 lucid/nn/modules/vision.py,sha256=8xYasT7TNj4NXwMwwJIw1nbV1paeWEFg_ZohXn9kZBg,1579
 lucid/optim/__init__.py,sha256=21EcCCPwrhPGP9TXvDje075_S2hPr0pHToygCaq8keI,201
-lucid/optim/_base.py,sha256=yHCjOw9AotqDl64nJyvYl6lVwGo4MpgPHwH_YAF9pKw,2022
+lucid/optim/_base.py,sha256=KxM5h5ONeO8hCpAzD2_vverFRKeymu2XC6AHN_L_v3g,4859
 lucid/optim/ada.py,sha256=POIl7dbv3qqwKxGGaceSrs-lZF1tD-vyvDxjtZdx--E,5807
 lucid/optim/adam.py,sha256=pVlZIcXD1s-IYK-WAfFognId8RhxzmlS5227-i0Vhq4,10347
 lucid/optim/prop.py,sha256=CbsWmoBb_g_8z16M3T6dMoSR9c72hm8M375IT1UHjpw,4740
 lucid/optim/sgd.py,sha256=DBZ1ZXQ9TfKZCRECfNRMDH9mvqUWCOPdY5TobnVxpz8,4477
 lucid/optim/lr_scheduler/__init__.py,sha256=kUoyN2g9nwTtEAqEVij832WSRvzEpKZywSJdfD7MQvY,58
-lucid/optim/lr_scheduler/_base.py,sha256=Qvw0OCaZQJaAmgkZbfqqYBb1sSKx0rE2obGHhN4u35E,2019
+lucid/optim/lr_scheduler/_base.py,sha256=NNJnjwmJpsRXathrbLtH4tjfBHtwOiJ5HwF1_S6Ym5c,3092
 lucid/optim/lr_scheduler/_schedulers.py,sha256=OIzduTXV6Ut4qcvw6neMPr3jlv6BgTSsys0-6KoHxK4,8140
 lucid/random/__init__.py,sha256=s8EAaKhEiTKT_vYjP4IFHx0xQVa1jqc_qIyvMauUu7M,2727
 lucid/random/_func.py,sha256=1Lu4m-ciEK037chNDGqv_j00RgGGzQ7UfslSfYActUk,2232
@@ -119,8 +120,8 @@ lucid/visual/__init__.py,sha256=6TuFDfmXTwpLyHl7_KqBfdzW6zqHjGzIFvymjFPlvjI,21
 lucid/visual/graph.py,sha256=YjpIDM_lloZARw3sCBiXPl_hT5A2gTk2fEHvwvJWXTk,4599
 lucid/weights/__init__.py,sha256=z1AikA3rOEeckWGkYWlcZkxNlJo9Xwa39PL6ly3hWnc,8801
 lucid/weights/__init__.pyi,sha256=lFonYC3cUx2Idolf3AEPnjFcyqcn3UDU84oJlZafqLY,3013
-lucid_dl-2.7.5.dist-info/licenses/LICENSE,sha256=vxRFYnVD1IeYtsvw-KmoElfqrjxKHv1h9YTvsG54loQ,1065
-lucid_dl-2.7.5.dist-info/METADATA,sha256=TFYQJ8ntSJLvF7J3qoAWsR7K95-DtnI-uprRQbQWzvM,11260
-lucid_dl-2.7.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lucid_dl-2.7.5.dist-info/top_level.txt,sha256=uzP_qBx9iNWIHKJRlElYcBLYVqMpdm9Q1Ma63QPYbFc,6
-lucid_dl-2.7.5.dist-info/RECORD,,
+lucid_dl-2.7.7.dist-info/licenses/LICENSE,sha256=vxRFYnVD1IeYtsvw-KmoElfqrjxKHv1h9YTvsG54loQ,1065
+lucid_dl-2.7.7.dist-info/METADATA,sha256=UVi3p93MdqSoYHvUkdil5C12wx5sfZy5Obw4Hkd7pxs,11260
+lucid_dl-2.7.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lucid_dl-2.7.7.dist-info/top_level.txt,sha256=uzP_qBx9iNWIHKJRlElYcBLYVqMpdm9Q1Ma63QPYbFc,6
+lucid_dl-2.7.7.dist-info/RECORD,,