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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

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

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

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)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

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)

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

# 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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

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

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

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()

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

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