module:设计

Author: miaobyte

front/py/deepx/nn/modules/module.py

@@ -10,9 +10,7 @@ def __init__(self, name: Optional[str] = None):
         self._parent: Optional[Module] = None
         self._modules: OrderedDict[str, Module] = OrderedDict()
         self._parameters: OrderedDict[str, Tensor] = OrderedDict()
-        self._buffers: OrderedDict[str, Tensor] = OrderedDict()
-        self.training = True
-        
+
     def _generate_default_name(self) -> str:
         class_name = self.__class__.__name__
         base_name = re.sub(r'(?<!^)(?=[A-Z])', '_', class_name).lower()
@@ -26,10 +24,7 @@ def __setattr__(self, name: str, value: Any) -> None:
         if isinstance(value, Module):
             self.register_module(name, value)
         elif isinstance(value, Tensor):
-            if value.requires_grad:
-                self.register_parameter(name, value)
-            else:
-                self.register_buffer(name, value)
+            self.register_parameter(name, value)
         super().__setattr__(name, value)
 
     def register_module(self, name: str, module: Optional['Module']) -> None:
@@ -46,14 +41,7 @@ def register_parameter(self, name: str, param: Optional[Tensor]) -> None:
         else:
             self._parameters[name] = param
             param.name = self._full_name + '.' + name
-            
-    def register_buffer(self, name: str, tensor: Optional[Tensor]) -> None:
-        if tensor is None:
-            self._buffers.pop(name, None)
-        else:
-            self._buffers[name] = tensor
-            tensor.name = self._full_name + '.' + name
-    
+
     @property
     def _full_name(self) -> str:
         names = []
@@ -123,19 +111,14 @@ def state_dict(self) -> Dict[str, Tensor]:
         state = {}
         for name, param in self.named_parameters():
             state[name] = param.detach().clone()
-        for name, buf in self.named_buffers():
-            state[name] = buf.detach().clone()
         return state
 
     def load_state_dict(self, state_dict: Dict[str, Tensor]) -> None:
         """加载模型状态"""
         for name, param in self.named_parameters():
             if name in state_dict:
                 param.data.copy_(state_dict[name])
-        for name, buf in self.named_buffers():
-            if name in state_dict:
-                buf.data.copy_(state_dict[name])
-                
+     
     def __call__(self, *args, **kwargs) -> Any:
         """允许模块像函数一样调用"""
         return self.forward(*args, **kwargs)