PyTorch神经网络模块系统深度剖析：从Module基类到复杂网络的构建机制

┌─────────────────────────────────────────┐
│         High-level Networks             │  ← 预定义网络结构
├─────────────────────────────────────────┤
│         Composite Modules               │  ← 组合模块(Sequential, ModuleList)
├─────────────────────────────────────────┤
│         Basic Layers                    │  ← 基础层(Linear, Conv, etc.)
├─────────────────────────────────────────┤
│         Module Base Class               │  ← Module基类
├─────────────────────────────────────────┤
│         Parameter & Buffer              │  ← 参数和缓冲区管理
├─────────────────────────────────────────┤
│         Hook System                     │  ← 钩子机制
└─────────────────────────────────────────┘

graph TB
    subgraph "PyTorch nn 模块系统架构"
        subgraph "用户接口层"
            NN_API[torch.nn API]
            FUNC[torch.nn.functional]
            INIT[torch.nn.init]
            UTILS[torch.nn.utils]
        end
        
        subgraph "高级网络构造"
            SEQ[Sequential]
            LIST[ModuleList]
            DICT[ModuleDict]
            CONTAINER[Container Modules]
        end
        
        subgraph "基础层实现"
            LINEAR[Linear Layer]
            CONV[Convolution Layers]
            NORM[Normalization Layers]
            ACT[Activation Layers]
            POOL[Pooling Layers]
            RNN[Recurrent Layers]
            LOSS[Loss Functions]
        end
        
        subgraph "Module基类系统"
            MODULE[Module Base]
            PARAM[Parameter]
            BUFFER[Buffer]
            STATE[State Management]
        end
        
        subgraph "钩子与回调系统"
            FORWARD_HOOK[Forward Hooks]
            BACKWARD_HOOK[Backward Hooks]
            PRE_HOOK[Pre Hooks]
            REGISTRATION[Registration Hooks]
        end
        
        subgraph "支持基础设施"
            DEVICE[Device Management]
            DTYPE[Data Type]
            TRAINING[Training Mode]
            GRAD[Gradient Management]
        end
        
        subgraph "底层支持"
            AUTOGRAD[Autograd System]
            TENSOR[Tensor Operations]
            FUNCTIONAL[Functional Backend]
            CUDA_IMPL[CUDA Implementation]
        end
    end
    
    %% 连接关系
    NN_API --> SEQ
    NN_API --> LIST
    NN_API --> DICT
    
    SEQ --> LINEAR
    SEQ --> CONV
    SEQ --> NORM
    
    LINEAR --> MODULE
    CONV --> MODULE
    NORM --> MODULE
    ACT --> MODULE
    POOL --> MODULE
    RNN --> MODULE
    LOSS --> MODULE
    
    MODULE --> PARAM
    MODULE --> BUFFER
    MODULE --> STATE
    
    MODULE --> FORWARD_HOOK
    MODULE --> BACKWARD_HOOK
    MODULE --> PRE_HOOK
    PARAM --> REGISTRATION
    
    MODULE --> DEVICE
    MODULE --> DTYPE
    MODULE --> TRAINING
    PARAM --> GRAD
    
    MODULE --> AUTOGRAD
    PARAM --> TENSOR
    FUNC --> FUNCTIONAL
    DEVICE --> CUDA_IMPL
    
    style MODULE fill:#e1f5fe
    style PARAM fill:#f3e5f5
    style FORWARD_HOOK fill:#e8f5e8
    style AUTOGRAD fill:#fff3e0

class Module:
    """神经网络模块基类的核心实现"""
    
    def __init__(self):
        # 训练模式标志
        self.training: bool = True
        
        # 参数字典 - 存储可学习参数
        self._parameters: Dict[str, Optional[Parameter]] = OrderedDict()
        
        # 缓冲区字典 - 存储非可学习但需要保存的张量
        self._buffers: Dict[str, Optional[Tensor]] = OrderedDict()
        
        # 非持久化缓冲区 - 不会被保存到state_dict
        self._non_persistent_buffers_set: Set[str] = set()
        
        # 子模块字典
        self._modules: Dict[str, Optional['Module']] = OrderedDict()
        
        # 钩子字典
        self._forward_hooks: Dict[int, Callable] = OrderedDict()
        self._forward_hooks_with_kwargs: Dict[int, bool] = OrderedDict()
        self._forward_hooks_always_called: Dict[int, bool] = OrderedDict()
        self._forward_pre_hooks: Dict[int, Callable] = OrderedDict()
        self._forward_pre_hooks_with_kwargs: Dict[int, bool] = OrderedDict()
        self._backward_hooks: Dict[int, Callable] = OrderedDict()
        self._backward_pre_hooks: Dict[int, Callable] = OrderedDict()
        
        # 状态管理
        self._state_dict_hooks: Dict[int, Callable] = OrderedDict()
        self._load_state_dict_pre_hooks: Dict[int, Callable] = OrderedDict()
        self._load_state_dict_post_hooks: Dict[int, Callable] = OrderedDict()
        
        # 版本控制
        self._version: int = 1

def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
    """注册可学习参数
    
    Args:
        name: 参数名称
        param: Parameter对象或None
    """
    if '_parameters' not in self.__dict__:
        raise AttributeError("cannot assign parameters before Module.__init__() call")
    
    elif not isinstance(name, str):
        raise TypeError(f"parameter name should be a string. Got {type(name)}")
    elif '.' in name:
        raise KeyError("parameter name can't contain '.'")
    elif name == '':
        raise KeyError("parameter name can't be empty string")
    elif hasattr(self, name) and name not in self._parameters:
        raise KeyError(f"attribute '{name}' already exists")
    
    # 如果参数为None，移除现有参数
    if param is None:
        self._parameters[name] = None
    elif not isinstance(param, Parameter):
        raise TypeError(f"cannot assign '{type(param)}' object to parameter '{name}' "
                       "(torch.nn.Parameter or None required)")
    elif param.grad_fn is not None:
        raise ValueError(f"cannot assign non-leaf Tensor to parameter '{name}'. Model "
                        f"parameters must be created explicitly. To express 'param + x' as a "
                        f"parameter, compute x outside the forward() method and use 'param = "
                        f"Parameter(param + x)'.")
    else:
        # 注册参数
        self._parameters[name] = param
    
    # 触发参数注册钩子
    for hook in _global_parameter_registration_hooks.values():
        output = hook(self, name, param)
        if output is not None:
            self._parameters[name] = output

def register_buffer(self, name: str, tensor: Optional[Tensor], persistent: bool = True) -> None:
    """注册缓冲区张量
    
    Args:
        name: 缓冲区名称
        tensor: 张量数据或None
        persistent: 是否持久化保存
    """
    if '_buffers' not in self.__dict__:
        raise AttributeError("cannot assign buffer before Module.__init__() call")
    elif not isinstance(name, str):
        raise TypeError(f"buffer name should be a string. Got {type(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(f"attribute '{name}' already exists")
    
    if tensor is not None and not isinstance(tensor, torch.Tensor):
        raise TypeError(f"cannot assign '{type(tensor)}' object to buffer '{name}' "
                       "(torch.Tensor or None required)")
    
    self._buffers[name] = tensor
    
    if not persistent:
        self._non_persistent_buffers_set.add(name)
    elif name in self._non_persistent_buffers_set:
        self._non_persistent_buffers_set.remove(name)
    
    # 触发缓冲区注册钩子
    for hook in _global_buffer_registration_hooks.values():
        output = hook(self, name, tensor)
        if output is not None:
            self._buffers[name] = output

def register_module(self, name: str, module: Optional['Module']) -> None:
    """注册子模块
    
    Args:
        name: 子模块名称
        module: Module对象或None
    """
    if not isinstance(name, str):
        raise TypeError(f"module name should be a string. Got {type(name)}")
    elif '.' in name:
        raise KeyError("module name can't contain '.'")
    elif name == '':
        raise KeyError("module name can't be empty string")
    elif hasattr(self, name) and name not in self._modules:
        raise KeyError(f"attribute '{name}' already exists")
    
    if module is not None and not isinstance(module, Module):
        raise TypeError(f"cannot assign '{type(module)}' object to child module '{name}' "
                       "(torch.nn.Module or None required)")
    
    self._modules[name] = module
    
    # 触发模块注册钩子
    for hook in _global_module_registration_hooks.values():
        output = hook(self, name, module)
        if output is not None:
            self._modules[name] = output

def __setattr__(self, name: str, value: Any) -> None:
    """设置属性时的拦截处理"""
    
    def remove_from(*dicts_or_sets):
        """从字典或集合中移除指定键"""
        for d in dicts_or_sets:
            if name in d:
                if isinstance(d, dict):
                    del d[name]
                else:
                    d.discard(name)
    
    params = self.__dict__.get('_parameters')
    if isinstance(value, Parameter):
        # 设置Parameter时，从其他容器中移除并注册为参数
        if params is None:
            raise AttributeError("cannot assign parameters before Module.__init__() call")
        remove_from(self.__dict__, self._buffers, self._modules, self._non_persistent_buffers_set)
        self.register_parameter(name, value)
        
    elif params is not None and name in params:
        # 如果已经存在同名参数，更新参数值
        if value is not None:
            raise TypeError(f"cannot assign '{type(value)}' as parameter '{name}' "
                           "(torch.nn.Parameter or None expected)")
        self.register_parameter(name, value)
        
    else:
        # 处理模块、缓冲区或普通属性
        modules = self.__dict__.get('_modules')
        if isinstance(value, Module):
            if modules is None:
                raise AttributeError("cannot assign module before Module.__init__() call")
            remove_from(self.__dict__, self._parameters, self._buffers, self._non_persistent_buffers_set)
            self.register_module(name, value)
            
        elif modules is not None and name in modules:
            if value is not None:
                raise TypeError(f"cannot assign '{type(value)}' as child module '{name}' "
                               "(torch.nn.Module or None expected)")
            self.register_module(name, value)
            
        else:
            # 处理缓冲区或普通属性
            buffers = self.__dict__.get('_buffers')
            if buffers is not None and name in buffers:
                if value is not None and not isinstance(value, torch.Tensor):
                    raise TypeError(f"cannot assign '{type(value)}' as buffer '{name}' "
                                   "(torch.Tensor or None expected)")
                self.register_buffer(name, value)
            else:
                # 普通属性赋值
                super().__setattr__(name, value)

def __getattr__(self, name: str) -> Any:
    """获取属性时的拦截处理"""
    if '_parameters' in self.__dict__:
        _parameters = self.__dict__['_parameters']
        if name in _parameters:
            return _parameters[name]
    if '_buffers' in self.__dict__:
        _buffers = self.__dict__['_buffers']
        if name in _buffers:
            return _buffers[name]
    if '_modules' in self.__dict__:
        modules = self.__dict__['_modules']
        if name in modules:
            return modules[name]
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

def __delattr__(self, name):
    """删除属性时的拦截处理"""
    if name in self._parameters:
        del self._parameters[name]
    elif name in self._buffers:
        del self._buffers[name]
        self._non_persistent_buffers_set.discard(name)
    elif name in self._modules:
        del self._modules[name]
    else:
        super().__delattr__(name)

class Parameter(torch.Tensor, metaclass=_ParameterMeta):
    """神经网络参数类
    
    Parameter是Tensor的子类，当赋值给Module属性时会自动被识别为参数
    """
    
    def __new__(cls, data=None, requires_grad=True):
        if data is None:
            data = torch.empty(0)
        
        if type(data) is torch.Tensor or type(data) is Parameter:
            # 标准Tensor路径：使用_make_subclass创建
            return torch.Tensor._make_subclass(cls, data, requires_grad)
        
        # 自定义张量路径：设置_is_param标志
        t = data.detach().requires_grad_(requires_grad)
        if type(t) is not type(data):
            raise RuntimeError(
                f"Creating a Parameter from an instance of type {type(data).__name__} "
                "requires that detach() returns an instance of the same type, but return "
                f"type {type(t).__name__} was found instead."
            )
        
        # 标记为参数
        t._is_param = True
        return t
    
    def __deepcopy__(self, memo):
        """深拷贝实现"""
        if id(self) in memo:
            return memo[id(self)]
        else:
            result = type(self)(
                self.data.clone(memory_format=torch.preserve_format), 
                self.requires_grad
            )
            memo[id(self)] = result
            return result
    
    def __repr__(self):
        """字符串表示"""
        return "Parameter containing:\n" + super().__repr__()
    
    def __reduce_ex__(self, proto):
        """序列化支持"""
        state = torch._utils._get_obj_state(self)
        
        # 不序列化钩子（见Note [Don't serialize hooks]）
        hooks = OrderedDict()
        if not state:
            return (
                torch._utils._rebuild_parameter,
                (self.data, self.requires_grad, hooks),
            )
        
        return (
            torch._utils._rebuild_parameter_with_state,
            (self.data, self.requires_grad, hooks, state),
        )

# Parameter元类，支持isinstance检查
class _ParameterMeta(torch._C._TensorMeta):
    def __instancecheck__(self, instance):
        """自定义isinstance行为"""
        if self is Parameter:
            if isinstance(instance, torch.Tensor) and getattr(instance, '_is_param', False):
                return True
        return super().__instancecheck__(instance)

# torch.nn.init模块的核心实现
import math
import warnings
from typing import Optional

def calculate_gain(nonlinearity, param=None):
    """计算非线性激活函数的增益
    
    Args:
        nonlinearity: 激活函数名称
        param: 激活函数参数（如leaky_relu的负斜率）
    
    Returns:
        推荐的增益值
    """
    linear_fns = ['linear', 'conv1d', 'conv2d', 'conv3d', 'conv_transpose1d', 
                  'conv_transpose2d', 'conv_transpose3d']
    
    if nonlinearity in linear_fns or nonlinearity == 'sigmoid':
        return 1
    elif nonlinearity == 'tanh':
        return 5.0 / 3
    elif nonlinearity == 'relu':
        return math.sqrt(2.0)
    elif nonlinearity == 'leaky_relu':
        if param is None:
            negative_slope = 0.01
        elif not isinstance(param, bool) and isinstance(param, int) or isinstance(param, float):
            negative_slope = param
        else:
            raise ValueError(f"negative_slope {param} not a valid number")
        return math.sqrt(2.0 / (1 + negative_slope ** 2))
    elif nonlinearity == 'selu':
        return 3.0 / 4
    else:
        raise ValueError(f"Unsupported nonlinearity {nonlinearity}")

def uniform_(tensor: Tensor, a: float = 0., b: float = 1.) -> Tensor:
    """用均匀分布初始化张量
    
    Args:
        tensor: n维张量
        a: 均匀分布下界
        b: 均匀分布上界
    """
    with torch.no_grad():
        tensor.uniform_(a, b)
    return tensor

def normal_(tensor: Tensor, mean: float = 0., std: float = 1.) -> Tensor:
    """用正态分布初始化张量"""
    with torch.no_grad():
        tensor.normal_(mean, std)
    return tensor

def xavier_uniform_(tensor: Tensor, gain: float = 1.) -> Tensor:
    """Xavier均匀分布初始化
    
    Xavier初始化保持前向传播时激活值的方差稳定
    """
    num_input = tensor.size(-1)  # 输入维度
    num_output = tensor.size(0)  # 输出维度
    
    # 计算均匀分布边界
    std = gain * math.sqrt(2.0 / float(num_input + num_output))
    a = math.sqrt(3.0) * std  # uniform分布的边界
    
    with torch.no_grad():
        tensor.uniform_(-a, a)
    return tensor

def xavier_normal_(tensor: Tensor, gain: float = 1.) -> Tensor:
    """Xavier正态分布初始化"""
    num_input = tensor.size(-1)
    num_output = tensor.size(0)
    
    std = gain * math.sqrt(2.0 / float(num_input + num_output))
    with torch.no_grad():
        tensor.normal_(0., std)
    return tensor

def kaiming_uniform_(tensor: Tensor, a: float = 0, mode: str = 'fan_in', 
                    nonlinearity: str = 'leaky_relu') -> Tensor:
    """Kaiming均匀分布初始化（He初始化）
    
    专门为ReLU激活函数设计的初始化方法
    """
    if 0 in tensor.shape:
        warnings.warn("Initializing zero-element tensors is a no-op")
        return tensor
    
    num_input, num_output = _calculate_fan_in_and_fan_out(tensor)
    
    if mode == 'fan_in':
        num = num_input
    elif mode == 'fan_out':
        num = num_output
    elif mode == 'fan_avg':
        num = (num_input + num_output) / 2
    else:
        raise ValueError(f"Mode {mode} not supported, please use one of "
                        "fan_in, fan_out or fan_avg.")
    
    gain = calculate_gain(nonlinearity, a)
    std = gain / math.sqrt(num)
    bound = math.sqrt(3.0) * std
    
    with torch.no_grad():
        tensor.uniform_(-bound, bound)
    return tensor

def _calculate_fan_in_and_fan_out(tensor):
    """计算扇入和扇出"""
    dimensions = tensor.dim()
    if dimensions < 2:
        raise ValueError("Fan in and fan out can not be computed for tensor with fewer than 2 dimensions")
    
    num_input_fmaps = tensor.size(1)
    num_output_fmaps = tensor.size(0)
    receptive_field_size = 1
    
    if tensor.dim() > 2:
        # 对于卷积核，计算感受野大小
        for s in tensor.shape[2:]:
            receptive_field_size *= s
    
    fan_in = num_input_fmaps * receptive_field_size
    fan_out = num_output_fmaps * receptive_field_size
    
    return fan_in, fan_out

def __call__(self, *args, **kwargs):
    """模块调用的主入口
    
    处理钩子执行、前向传播调用和结果后处理
    """
    # 检查是否需要全局前向钩子
    for hook_id, hook in _global_forward_pre_hooks.items():
        if hook_id in _global_forward_hooks_always_called or _has_any_global_hook():
            result = hook(self, args, kwargs)
            if result is not None:
                if isinstance(result, tuple) and len(result) == 2:
                    args, kwargs = result
                else:
                    args = result
    
    # 执行模块级前向预钩子
    bw_hook = None
    if len(self._forward_pre_hooks) > 0:
        for hook_id, hook in self._forward_pre_hooks.items():
            if hook_id in self._forward_pre_hooks_with_kwargs:
                hook_result = hook(self, args, kwargs)
            else:
                hook_result = hook(self, args)
            
            if hook_result is not None:
                if isinstance(hook_result, tuple) and len(hook_result) == 2:
                    args, kwargs = hook_result
                else:
                    args = hook_result
    
    # 核心前向传播调用
    result = self.forward(*args, **kwargs)
    
    # 执行前向后钩子
    for hook_id, hook in self._forward_hooks.items():
        if hook_id in self._forward_hooks_with_kwargs:
            hook_result = hook(self, args, kwargs, result)
        else:
            hook_result = hook(self, args, result)
        
        if hook_result is not None:
            result = hook_result
    
    # 执行全局前向后钩子
    for hook_id, hook in _global_forward_hooks.items():
        if hook_id in _global_forward_hooks_with_kwargs:
            hook_result = hook(self, args, kwargs, result)
        else:
            hook_result = hook(self, args, result)
        
        if hook_result is not None:
            result = hook_result
    
    # 如果有反向钩子需要注册
    if bw_hook:
        result.register_hook(bw_hook)
    
    return result

def forward(self, *input):
    """前向传播的抽象方法
    
    子类必须重写此方法来定义具体的前向传播逻辑
    """
    raise NotImplementedError

class Linear(Module):
    """线性层实现"""
    
    __constants__ = ['in_features', 'out_features']
    in_features: int
    out_features: int
    weight: Tensor
    
    def __init__(self, in_features: int, out_features: int, bias: bool = True, 
                 device=None, dtype=None) -> None:
        """初始化线性层
        
        Args:
            in_features: 输入特征数
            out_features: 输出特征数  
            bias: 是否使用偏置
        """
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        
        self.in_features = in_features
        self.out_features = out_features
        
        # 创建权重参数 [out_features, in_features]
        self.weight = Parameter(torch.empty((out_features, in_features), **factory_kwargs))
        
        if bias:
            # 创建偏置参数 [out_features]
            self.bias = Parameter(torch.empty(out_features, **factory_kwargs))
        else:
            # 注册为None，表示不使用偏置
            self.register_parameter('bias', None)
        
        # 初始化参数
        self.reset_parameters()
    
    def reset_parameters(self) -> None:
        """重置参数为默认初始值"""
        # 使用Kaiming均匀分布初始化权重
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        
        if self.bias is not None:
            # 偏置使用小范围均匀分布初始化
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            init.uniform_(self.bias, -bound, bound)
    
    def forward(self, input: Tensor) -> Tensor:
        """前向传播：y = xW^T + b
        
        Args:
            input: 输入张量 [*, in_features]
            
        Returns:
            输出张量 [*, out_features]
        """
        return F.linear(input, self.weight, self.bias)
    
    def extra_repr(self) -> str:
        """返回层的额外信息用于打印"""
        return f'in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}'

# torch.nn.functional.linear的实现
def linear(input: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
    """线性变换的函数式实现
    
    Args:
        input: 输入张量 [..., in_features]
        weight: 权重矩阵 [out_features, in_features]
        bias: 偏置向量 [out_features]，可选
    
    Returns:
        输出张量 [..., out_features]
    """
    if input.dim() == 2 and bias is not None:
        # 优化的2D情况：使用addmm
        ret = torch.addmm(bias, input, weight.t())
    else:
        # 通用情况：使用矩阵乘法
        output = input.matmul(weight.t())
        if bias is not None:
            output += bias
        ret = output
    
    return ret

class Module:
    """钩子系统的完整实现"""
    
    def register_forward_pre_hook(
        self, 
        hook: Callable[..., None], 
        *,
        prepend: bool = False,
        with_kwargs: bool = False
    ) -> RemovableHandle:
        """注册前向传播前钩子
        
        钩子签名：
        - hook(module, input) -> None or modified input
        - hook(module, input, kwargs) -> None or (modified input, kwargs) (如果with_kwargs=True)
        
        Args:
            hook: 钩子函数
            prepend: 是否添加到钩子列表开头
            with_kwargs: 是否传递kwargs参数
            
        Returns:
            可移除的钩子句柄
        """
        handle = RemovableHandle(self._forward_pre_hooks)
        self._forward_pre_hooks[handle.id] = hook
        
        if with_kwargs:
            self._forward_pre_hooks_with_kwargs[handle.id] = True
        
        if prepend:
            # 将钩子移动到开头
            self._forward_pre_hooks.move_to_end(handle.id, last=False)
        
        return handle
    
    def register_forward_hook(
        self,
        hook: Callable[..., None],
        *,
        prepend: bool = False,
        with_kwargs: bool = False,
        always_call: bool = False
    ) -> RemovableHandle:
        """注册前向传播后钩子
        
        钩子签名：
        - hook(module, input, output) -> None or modified output
        - hook(module, input, kwargs, output) -> None or modified output (如果with_kwargs=True)
        
        Args:
            hook: 钩子函数
            prepend: 是否添加到钩子列表开头
            with_kwargs: 是否传递kwargs参数
            always_call: 是否总是调用（即使模块不在训练模式）
        """
        handle = RemovableHandle(self._forward_hooks)
        self._forward_hooks[handle.id] = hook
        
        if with_kwargs:
            self._forward_hooks_with_kwargs[handle.id] = True
        if always_call:
            self._forward_hooks_always_called[handle.id] = True
        
        if prepend:
            self._forward_hooks.move_to_end(handle.id, last=False)
        
        return handle
    
    def register_full_backward_pre_hook(
        self,
        hook: Callable[..., None],
        prepend: bool = False
    ) -> RemovableHandle:
        """注册完整的反向传播前钩子
        
        钩子签名：hook(module, grad_output) -> None or modified grad_output
        """
        handle = RemovableHandle(self._backward_pre_hooks)
        self._backward_pre_hooks[handle.id] = hook
        
        if prepend:
            self._backward_pre_hooks.move_to_end(handle.id, last=False)
        
        return handle
    
    def register_full_backward_hook(
        self,
        hook: Callable[..., None],
        prepend: bool = False
    ) -> RemovableHandle:
        """注册完整的反向传播后钩子
        
        钩子签名：hook(module, grad_input, grad_output) -> None or modified grad_input
        """
        handle = RemovableHandle(self._backward_hooks)
        self._backward_hooks[handle.id] = hook
        
        if prepend:
            self._backward_hooks.move_to_end(handle.id, last=False)
        
        return handle

# 全局钩子注册函数
def register_module_forward_pre_hook(hook: Callable[..., None]) -> RemovableHandle:
    """注册全局前向传播前钩子"""
    handle = RemovableHandle(_global_forward_pre_hooks)
    _global_forward_pre_hooks[handle.id] = hook
    return handle

def register_module_forward_hook(hook: Callable[..., None]) -> RemovableHandle:
    """注册全局前向传播后钩子"""
    handle = RemovableHandle(_global_forward_hooks)
    _global_forward_hooks[handle.id] = hook
    return handle

import torch
import torch.nn as nn

class HookExample:
    """钩子使用示例"""
    
    def __init__(self):
        self.activations = {}
        self.gradients = {}
    
    def activation_hook(self, name):
        """创建激活值钩子"""
        def hook(module, input, output):
            # 保存激活值
            self.activations[name] = output.detach().clone()
            print(f"Forward hook for {name}: output shape {output.shape}")
        return hook
    
    def gradient_hook(self, name):
        """创建梯度钩子"""
        def hook(module, grad_input, grad_output):
            # 保存梯度
            if grad_output[0] is not None:
                self.gradients[name] = grad_output[0].detach().clone()
                print(f"Backward hook for {name}: grad_output shape {grad_output[0].shape}")
        return hook
    
    def register_hooks(self, model):
        """为模型注册钩子"""
        hooks = []
        
        for name, module in model.named_modules():
            if isinstance(module, (nn.Linear, nn.Conv2d)):
                # 注册前向钩子
                handle1 = module.register_forward_hook(self.activation_hook(name))
                hooks.append(handle1)
                
                # 注册反向钩子
                handle2 = module.register_full_backward_hook(self.gradient_hook(name))
                hooks.append(handle2)
        
        return hooks

# 使用示例
def hook_usage_example():
    # 创建简单模型
    model = nn.Sequential(
        nn.Linear(10, 20),
        nn.ReLU(),
        nn.Linear(20, 1)
    )
    
    # 注册钩子
    hook_manager = HookExample()
    handles = hook_manager.register_hooks(model)
    
    # 前向传播
    x = torch.randn(5, 10, requires_grad=True)
    output = model(x)
    loss = output.sum()
    
    # 反向传播
    loss.backward()
    
    # 查看保存的激活值和梯度
    print("Saved activations:", list(hook_manager.activations.keys()))
    print("Saved gradients:", list(hook_manager.gradients.keys()))
    
    # 移除钩子
    for handle in handles:
        handle.remove()

def state_dict(self, *args, destination=None, prefix='', keep_vars=False):
    """获取模块状态字典
    
    Args:
        destination: 目标字典，如果为None则创建新的
        prefix: 键名前缀
        keep_vars: 是否保持Variable类型
        
    Returns:
        包含模块状态的字典
    """
    if destination is None:
        destination = OrderedDict()
    
    # 保存参数
    for name, param in self._parameters.items():
        if param is not None:
            key = prefix + name
            destination[key] = param if keep_vars else param.detach()
    
    # 保存持久化缓冲区
    for name, buf in self._buffers.items():
        if buf is not None and name not in self._non_persistent_buffers_set:
            key = prefix + name
            destination[key] = buf if keep_vars else buf.detach()
    
    # 递归保存子模块状态
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(destination=destination, prefix=prefix + name + '.')
    
    # 执行状态字典钩子
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, {})
        if hook_result is not None:
            destination = hook_result
    
    return destination

def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
    """加载模块状态
    
    Args:
        state_dict: 状态字典
        strict: 是否严格模式（键必须完全匹配）
        assign: 是否使用赋值而非复制
        
    Returns:
        NamedTuple包含missing_keys和unexpected_keys
    """
    if not isinstance(state_dict, Mapping):
        raise TypeError(f"Expected state_dict to be dict-like, got {type(state_dict)}")
    
    missing_keys = []
    unexpected_keys = []
    error_msgs = []
    
    # 执行预加载钩子
    for hook in self._load_state_dict_pre_hooks.values():
        hook(state_dict, '', {}, strict, missing_keys, unexpected_keys, error_msgs)
    
    def load(module, prefix=''):
        # 获取模块的本地状态
        local_state = {k: v for k, v in module._parameters.items() if v is not None}
        local_state.update({k: v for k, v in module._buffers.items() 
                           if v is not None and k not in module._non_persistent_buffers_set})
        
        # 检查状态字典中的键
        for name, param in local_state.items():
            key = prefix + name
            if key in state_dict:
                input_param = state_dict[key]
                
                # 检查形状兼容性
                if not _same_shape(input_param, param):
                    error_msgs.append(f'size mismatch for {key}: copying a param with shape {input_param.shape} from checkpoint, '
                                     f'the shape in current model is {param.shape}.')
                    continue
                
                try:
                    if assign:
                        # 使用赋值
                        if isinstance(param, Parameter):
                            param.data = input_param
                        else:
                            param.copy_(input_param)
                    else:
                        # 使用复制
                        param.copy_(input_param)
                        
                except Exception as ex:
                    error_msgs.append(f'While copying the parameter named "{key}", '
                                     f'whose dimensions in the model are {param.shape} '
                                     f'and whose dimensions in the checkpoint are {input_param.shape}, '
                                     f'an exception occurred : {ex.args}.')
            elif strict:
                missing_keys.append(key)
        
        # 递归加载子模块
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')
    
    # 执行加载
    load(self)
    del load
    
    # 查找意外的键
    if strict:
        for key in state_dict.keys():
            if not any(key.startswith(prefix) for prefix in self._modules.keys()):
                unexpected_keys.append(key)
    
    # 执行后加载钩子
    for hook in self._load_state_dict_post_hooks.values():
        hook(self, '', {}, strict, missing_keys, unexpected_keys, error_msgs)
    
    if error_msgs:
        raise RuntimeError('Error(s) in loading state_dict:\n\t' + '\n\t'.join(error_msgs))
    
    return _IncompatibleKeys(missing_keys, unexpected_keys)

def _same_shape(tensor1, tensor2):
    """检查两个张量形状是否相同"""
    return tensor1.shape == tensor2.shape

def train(self: T, mode: bool = True) -> T:
    """设置训练模式
    
    Args:
        mode: True为训练模式，False为评估模式
        
    Returns:
        返回self以支持链式调用
    """
    if not isinstance(mode, bool):
        raise ValueError("training mode is expected to be boolean")
    
    self.training = mode
    
    # 递归设置所有子模块的训练模式
    for module in self.children():
        module.train(mode)
    
    return self

def eval(self: T) -> T:
    """设置为评估模式"""
    return self.train(False)

def requires_grad_(self: T, requires_grad: bool = True) -> T:
    """设置所有参数是否需要梯度
    
    Args:
        requires_grad: 是否需要梯度计算
        
    Returns:
        返回self以支持链式调用
    """
    for p in self.parameters():
        p.requires_grad_(requires_grad)
    return self

def to(self: T, *args, **kwargs) -> T:
    """将模块转移到指定设备/数据类型
    
    支持多种调用方式：
    - module.to(device)
    - module.to(dtype)
    - module.to(tensor)  # 推断设备和数据类型
    - module.to(device, dtype)
    """
    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
    
    if dtype is not None:
        if not (dtype.is_floating_point or dtype.is_complex):
            raise TypeError(f'nn.Module.to only accepts floating point or complex dtypes, but got desired dtype={dtype}')
    
    def convert(t):
        if convert_to_format is not None and t.dim() in (4, 5):
            return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
                       non_blocking, memory_format=convert_to_format)
        return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)
    
    return self._apply(convert)

def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:
    """将模块移动到CUDA设备"""
    return self._apply(lambda t: t.cuda(device))

def xpu(self: T, device: Optional[Union[int, device]] = None) -> T:
    """将模块移动到XPU设备"""
    return self._apply(lambda t: t.xpu(device))

def cpu(self: T) -> T:
    """将模块移动到CPU"""
    return self._apply(lambda t: t.cpu())

def type(self: T, dst_type: Union[dtype, str]) -> T:
    """转换模块的数据类型"""
    return self._apply(lambda t: t.type(dst_type))

def float(self: T) -> T:
    """转换为float类型"""
    return self._apply(lambda t: t.float() if t.is_floating_point() else t)

def double(self: T) -> T:
    """转换为double类型"""
    return self._apply(lambda t: t.double() if t.is_floating_point() else t)

def half(self: T) -> T:
    """转换为半精度类型"""
    return self._apply(lambda t: t.half() if t.is_floating_point() else t)

def bfloat16(self: T) -> T:
    """转换为bfloat16类型"""
    return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)

def _apply(self, fn):
    """应用函数到所有参数、缓冲区和子模块
    
    这是设备和类型转换的核心实现
    """
    for module in self.children():
        module._apply(fn)
    
    def compute_should_use_set_data(tensor, tensor_applied):
        """计算是否应该使用set_data而非copy_"""
        if torch._has_compatible_shallow_copy_type(tensor, tensor_applied):
            return False
        
        if tensor.dtype != tensor_applied.dtype:
            return True
        if tensor.device != tensor_applied.device:
            return True
        if tensor.is_sparse != tensor_applied.is_sparse:
            return True
        return False
    
    # 转换参数
    for key, param in self._parameters.items():
        if param is None:
            continue
        
        # 子类可以通过重写这个方法来定制转换行为
        with torch.no_grad():
            param_applied = fn(param)
            should_use_set_data = compute_should_use_set_data(param, param_applied)
            
            if should_use_set_data:
                param.data = param_applied
            else:
                param.copy_(param_applied)
    
    # 转换缓冲区
    for key, buf in self._buffers.items():
        if buf is not None:
            self._buffers[key] = fn(buf)
    
    return self

class Sequential(Module):
    """序列容器：按照添加顺序执行子模块"""
    
    def __init__(self, *args):
        super().__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict):
            # 从OrderedDict构造
            for key, module in args[0].items():
                self.add_module(key, module)
        else:
            # 从位置参数构造
            for idx, module in enumerate(args):
                self.add_module(str(idx), module)
    
    def _get_item_by_idx(self, iterator, idx) -> T:
        """通过索引获取模块"""
        size = len(self)
        idx = operator.index(idx)
        if not -size <= idx < size:
            raise IndexError(f'index {idx} is out of range')
        idx %= size
        return next(itertools.islice(iterator, idx, None))
    
    def __getitem__(self, idx) -> Union['Module', 'Sequential']:
        if isinstance(idx, slice):
            # 切片访问
            return self.__class__(OrderedDict(list(self._modules.items())[idx]))
        else:
            # 单个索引访问
            return self._get_item_by_idx(self._modules.values(), idx)
    
    def __setitem__(self, idx: int, module: Module) -> None:
        """设置指定索引的模块"""
        key: str = self._get_item_by_idx(self._modules.keys(), idx)
        return setattr(self, key, module)
    
    def __delitem__(self, idx: Union[slice, int]) -> None:
        """删除指定索引的模块"""
        if isinstance(idx, slice):
            for key in list(self._modules.keys())[idx]:
                delattr(self, key)
        else:
            key = self._get_item_by_idx(self._modules.keys(), idx)
            delattr(self, key)
        
        # 重新编号
        str_indices = [str(i) for i in range(len(self._modules))]
        self._modules = OrderedDict(list(zip(str_indices, self._modules.values())))
    
    def __len__(self) -> int:
        return len(self._modules)
    
    def __add__(self, other) -> 'Sequential':
        """连接两个Sequential"""
        if isinstance(other, Sequential):
            ret = Sequential()
            for layer in self:
                ret.append(layer)
            for layer in other:
                ret.append(layer)
            return ret
        else:
            raise ValueError('add operator supports only objects of Sequential class')
    
    def pop(self, key: Union[int, slice]) -> Module:
        """移除并返回指定模块"""
        v = self[key]
        del self[key]
        return v
    
    def append(self, module: Module) -> 'Sequential':
        """添加模块到末尾"""
        self.add_module(str(len(self)), module)
        return self
    
    def prepend(self, module: Module) -> 'Sequential':
        """添加模块到开头"""
        self._modules = OrderedDict([(str(0), module)] + 
                                   [(str(i + 1), m) for i, m in enumerate(self._modules.values())])
        return self
    
    def insert(self, index: int, module: Module) -> 'Sequential':
        """在指定位置插入模块"""
        if not isinstance(module, Module):
            raise AssertionError(f'module should be of type Module, got: {type(module)}')
        
        n = len(self._modules)
        if not (-n <= index <= n):
            raise IndexError(f'Index out of range: {index}')
        
        if index < 0:
            index += n
        
        for i in range(n, index, -1):
            self._modules[str(i)] = self._modules[str(i - 1)]
        
        self._modules[str(index)] = module
        return self
    
    def extend(self, sequential) -> 'Sequential':
        """扩展Sequential"""
        for layer in sequential:
            self.append(layer)
        return self
    
    def forward(self, input):
        """前向传播：依次通过所有子模块"""
        for module in self:
            input = module(input)
        return input