feat: add register buffer in InfiniCoreModule

Author: spike-zhu

python/infinicore/nn/modules/module.py

@@ -14,9 +14,7 @@
 from collections import OrderedDict, namedtuple
 import itertools
 import torch
-import torch.nn as nn
 from typing import Optional, Iterator, Tuple, Dict, Mapping, List, Any, overload, TypeVar
-import inspect
 import warnings
 
 _EXTRA_STATE_KEY_SUFFIX = '_extra_state'
@@ -95,6 +93,7 @@ def remove_from(*dicts_or_sets) -> None:
                     self.__dict__,
                     self._parameters,
                     self._buffers,
+                    self._non_persistent_buffers_set,
                 )
                 modules[name] = value
             elif modules is not None and name in modules:
@@ -113,7 +112,63 @@ def remove_from(*dicts_or_sets) -> None:
                                         )
                     buffers[name] = value
                 else:
-                    super.__setattr__(name, value)
+                    super().__setattr__(name, value)
+
+    def register_buffer(self, name: str, tensor: Optional[torch.tensor], persistent: bool = True) -> None:
+        r"""Adds a buffer to the module.
+
+        This is typically used to register a buffer that should not to be
+        considered a model parameter. For example, BatchNorm's ``running_mean``
+        is not a parameter, but is part of the module's state. Buffers, by
+        default, are persistent and will be saved alongside parameters. This
+        behavior can be changed by setting :attr:`persistent` to ``False``. The
+        only difference between a persistent buffer and a non-persistent buffer
+        is that the latter will not be a part of this module's
+        :attr:`state_dict`.
+
+        Buffers can be accessed as attributes using given names.
+
+        Args:
+            name (str): name of the buffer. The buffer can be accessed
+                from this module using the given name
+            tensor (Tensor or None): buffer to be registered. If ``None``, then operations
+                that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,
+                the buffer is **not** included in the module's :attr:`state_dict`.
+            persistent (bool): whether the buffer is part of this module's
+                :attr:`state_dict`.
+
+        Example::
+
+            >>> # xdoctest: +SKIP("undefined vars")
+            >>> self.register_buffer('running_mean', torch.zeros(num_features))
+
+        """
+        if persistent is False and isinstance(self, torch.jit.ScriptModule):
+            raise RuntimeError("ScriptModule does not support non-persistent buffers")
+
+        if '_buffers' not in self.__dict__:
+            raise AttributeError(
+                "cannot assign buffer before Module.__init__() call")
+        elif not isinstance(name, str):
+            raise TypeError("buffer name should be a string. "
+                            "Got {}".format(torch.typename(name)))
+        elif '.' in name:
+            raise KeyError("buffer name can't contain \".\"")
+        elif name == '':
+            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, torch.Tensor):
+            raise TypeError("cannot assign '{}' object to buffer '{}' "
+                            "(torch Tensor or None required)"
+                            .format(torch.typename(tensor), name))
+        else:
+            self._buffers[name] = tensor
+            if persistent:
+                self._non_persistent_buffers_set.discard(name)
+            else:
+                self._non_persistent_buffers_set.add(name)
+
 
     def register_parameter(self, name: str, param: Optional[torch.nn.Parameter]) -> None:
         r"""Add a parameter to the module.