Merge pull request #2 from pengcheng888/feature/add_nn_interface

Feature/add nn interface

Author: gongchensu

python/infinicore/nn/__init__.py

@@ -1,8 +1,5 @@
-from infinicore.nn import (
-    functional as functional,
-)
-from infinicore.nn import (
-    modules as modules,
-)
-from infinicore.nn.functional import *  # noqa: F403
+from infinicore.nn import functional
 from infinicore.nn.modules import *  # noqa: F403
+from infinicore.nn.parameter import InfiniCoreParameter as Parameter
+
+__all__ = ["functional", "Parameter"]

python/infinicore/nn/modules/__init__.py

@@ -1,5 +1,4 @@
-from .container import ModuleList
-from .module import Module
-from .parameter import Parameter
+from .container import InfiniCoreModuleList as ModuleList
+from .module import InfiniCoreModule as Module
 
-__all__ = ["ModuleList", "Module", "Parameter"]
+__all__ = ["ModuleList", "Module"]

python/infinicore/nn/modules/container.py

@@ -9,7 +9,7 @@
 from itertools import chain
 from typing import Iterator, List, Optional, Sequence, TypeVar, Union
 
-from .module import Module
+from .module import InfiniCoreModule as Module
 
 # Define type variable for module compatibility (supports InfiniCoreModule)
 ModuleType = TypeVar("ModuleType", bound=Union["Module"])
@@ -186,6 +186,3 @@ def __dir__(self) -> List[str]:
         # Filter out numeric keys to avoid cluttering dir() output
         keys = [key for key in keys if not key.isdigit()]
         return keys
-
-
-ModuleList = InfiniCoreModuleList

python/infinicore/nn/modules/module.py

@@ -32,7 +32,8 @@
 
 import infinicore
 
-from .parameter import Parameter
+from ...tensor import Tensor
+from ..parameter import InfiniCoreParameter as Parameter
 
 _EXTRA_STATE_KEY_SUFFIX = "_extra_state"
 T = TypeVar("T", bound="InfiniCoreModule")
@@ -57,7 +58,7 @@ class InfiniCoreModule:
 
     _version: int = 1
     _parameters: Dict[str, Optional[Parameter]]
-    _buffers: Dict[str, Optional[infinicore.Tensor]]
+    _buffers: Dict[str, Optional[Tensor]]
     _non_persistent_buffers_set: Set[str]
     _modules: Dict[str, Optional["InfiniCoreModule"]]
 
@@ -87,9 +88,7 @@ def __getattr__(self, name: str) -> Any:
             f"'{type(self).__name__}' object has no attribute '{name}'"
         )
 
-    def __setattr__(
-        self, name: str, value: Union[infinicore.Tensor, "InfiniCoreModule"]
-    ) -> None:
+    def __setattr__(self, name: str, value: Union[Tensor, "InfiniCoreModule"]) -> None:
         def remove_from(*dicts_or_sets) -> None:
             for d in dicts_or_sets:
                 if name in d:
@@ -113,7 +112,7 @@ def remove_from(*dicts_or_sets) -> None:
             )
             self.register_parameter(name, value)
         elif name in params:  # value will overwrite the name of params.
-            if not isinstance(value, infinicore.Tensor):
+            if not isinstance(value, Tensor):
                 raise TypeError(
                     f"cannot assign 'value' as parameter '{name}'  (infinicore.nn.Parameter, Parameter or None expected)"
                 )
@@ -141,7 +140,7 @@ def remove_from(*dicts_or_sets) -> None:
             else:
                 buffers = self.__dict__.get("_buffers")
                 if buffers is not None and name in buffers:
-                    if value is not None and not isinstance(value, infinicore.Tensor):
+                    if value is not None and not isinstance(value, Tensor):
                         raise TypeError(
                             f"cannot assign 'value' as buffer '{name}' "
                             "(torch.Tensor or None expected)"
@@ -154,7 +153,7 @@ def __call__(self, *input, **kwargs):
         return self.forward(*input, **kwargs)
 
     def register_buffer(
-        self, name: str, tensor: Optional[infinicore.Tensor], persistent: bool = True
+        self, name: str, tensor: Optional[Tensor], persistent: bool = True
     ) -> None:
         r"""Adds a buffer to the module.
 
@@ -187,7 +186,7 @@ def register_buffer(
             raise KeyError('buffer name can\'t be empty string ""')
         elif hasattr(self, name) and name not in self._buffers:
             raise KeyError("attribute '{}' already exists".format(name))
-        elif tensor is not None and not isinstance(tensor, infinicore.Tensor):
+        elif tensor is not None and not isinstance(tensor, Tensor):
             raise TypeError(
                 "cannot assign '{}' object to buffer '{}' "
                 "(torch Tensor or None required)".format("tensor", name)
@@ -256,7 +255,7 @@ def register_parameter(self, name: str, param: Parameter) -> None:
         if param is None:
             self._parameters[name] = None  # 竟然可以是None
         else:
-            if not isinstance(param, (Parameter, infinicore.Tensor)):
+            if not isinstance(param, (Parameter, Tensor)):
                 raise TypeError(
                     f"cannot assign  'param' object to parameter '{name}' "
                     "(infinicore.nn.Parameter, Parameter or None required)"
@@ -477,9 +476,9 @@ def _load_from_state_dict(
                 input_param = state_dict[key]
 
                 # input_param must be of type infinicore.Tensor
-                if not isinstance(input_param, infinicore.Tensor):
+                if not isinstance(input_param, Tensor):
                     raise TypeError(
-                        f"While copying the parameter named {key}, expected infinicore.Tensor from checkpoint but received {type(input_param)}"
+                        f"While copying the parameter named {key}, expected Tensor from checkpoint but received {type(input_param)}"
                     )
 
                 if (
@@ -575,7 +574,7 @@ def load(module, local_state_dict, prefix=""):
                         for k, v in local_state_dict.items()
                         if k.startswith(child_prefix)
                     }
-                    load(child, child_state_dict, child_prefix)  # noqa: F821
+                    load(child, child_state_dict, child_prefix)
 
         load(self, state_dict)
         del load
@@ -654,7 +653,7 @@ def named_parameters(
         for elem in gen:
             yield elem
 
-    def buffers(self, recurse: bool = True) -> Iterator[infinicore.Tensor]:
+    def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
         r"""Returns an iterator over module buffers.
 
         Args:
@@ -677,7 +676,7 @@ def buffers(self, recurse: bool = True) -> Iterator[infinicore.Tensor]:
 
     def named_buffers(
         self, prefix: str = "", recurse: bool = True
-    ) -> Iterator[Tuple[str, infinicore.Tensor]]:
+    ) -> Iterator[Tuple[str, Tensor]]:
         r"""Returns an iterator over module buffers, yielding both the
         name of the buffer as well as the buffer itself.
 
@@ -856,6 +855,3 @@ def _apply(self, fn, recurse=True):
 
     def to(self, *args, **kwargs):
         raise KeyError("not support")
-
-
-Module = InfiniCoreModule

python/infinicore/nn/modules/parameter.py renamed to python/infinicore/nn/parameter.py

@@ -13,14 +13,14 @@
 # The use of this file is governed by the BSD 3-Clause License.
 
 
-import infinicore
+from ..tensor import Tensor
 
 
-class InfiniCoreParameter(infinicore.Tensor):
+class InfiniCoreParameter(Tensor):
     r"""A kind of Tensor that is to be considered a module parameter."""
 
     def __init__(self, data=None):
-        if not isinstance(data, infinicore.Tensor):
+        if not isinstance(data, Tensor):
             raise ValueError("The `data` variable must be of type `infinicore.Tensor`.")
         super().__init__(data._underlying)
 
@@ -32,6 +32,3 @@ def __deepcopy__(self, memo):
 
     def __reduce_ex__(self, proto):
         raise ValueError("not supported!")
-
-
-Parameter = InfiniCoreParameter