如何获取PyTorch模型前向传播中实际调用的所有层（按顺序）

之前查找PyTorch获取模型所有层的方法时，看到的方案都是遍历children()或named_modules()，但用这些方法处理ResNet50时遇到了问题：ResNet的Bottleneck模块里只定义了一个ReLU层实例，但在forward函数中被调用了三次。现有方法只能解析出这一个ReLU层，无法反映前向传播中实际调用的三次ReLU操作，我需要能按前向传播顺序获取所有实际执行的层调用记录。

以下是ResNet Bottleneck的源码：

class Bottleneck(nn.Module):
    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
    # according to "Deep residual learning for image recognition"
    # This variant is also known as ResNet V1.5 and improves accuracy according to relevant NVIDIA docs.

    expansion: int = 4

    def __init__(
        self,
        inplanes: int,
        planes: int,
        stride: int = 1,
        downsample: Optional[nn.Module] = None,
        groups: int = 1,
        base_width: int = 64,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None,
    ) -> None:
        super().__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.0)) * groups
        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out

解决方案：使用前向钩子追踪层调用顺序 #

要获取前向传播中实际执行的每一层调用，我们可以利用PyTorch的**前向钩子（Forward Hook）**来记录每次层的执行。钩子会在每个模块执行前/后触发，借此按顺序记录所有被调用的模块实例，包括同一实例的多次调用。

实现步骤： #

• 定义钩子函数，用于记录每次模块调用
• 遍历模型所有模块，为每个模块注册前向钩子
• 输入测试张量进行前向传播，触发钩子记录顺序
• 整理输出记录的调用序列

代码实现： #

import torch
import torch.nn as nn
from torchvision.models import resnet50

# 存储调用顺序的列表
call_sequence = []

def forward_hook(module, input, output):
    # 记录模块类型和实例，可按需调整记录内容
    call_sequence.append((module.__class__.__name__, module))

def get_forward_call_sequence(model, input_tensor):
    global call_sequence
    call_sequence.clear()
    
    # 为非容器类模块注册钩子
    hooks = []
    for module in model.modules():
        if not isinstance(module, (nn.Sequential, nn.ModuleList)):
            hook = module.register_forward_hook(forward_hook)
            hooks.append(hook)
    
    # 前向传播触发钩子
    model(input_tensor)
    
    # 移除钩子避免后续影响
    for hook in hooks:
        hook.remove()
    
    return call_sequence

# 测试ResNet50
model = resnet50(pretrained=False)
input_tensor = torch.randn(1, 3, 224, 224)
sequence = get_forward_call_sequence(model, input_tensor)

# 输出前15次调用验证
print("前15次调用序列：")
for idx, (name, _) in enumerate(sequence[:15]):
    print(f"{idx+1}. {name}")

说明： #

• 钩子函数在模块执行后触发，记录顺序与前向传播执行顺序完全一致
• 对于Bottleneck中被多次调用的ReLU实例，每次调用都会被单独记录，最终序列里会出现三次该实例的记录
• 如果需要区分同一模块的不同调用位置，可以额外记录模块的名称路径（通过named_modules()获取）

替代方案：修改模型拆分复用层 #

如果允许修改模型源码，可以将Bottleneck中的单个ReLU拆分为三个独立实例，这样named_modules()就能直接获取到三个ReLU层，自然反映前向顺序：

修改后的Bottleneck __init__部分：

def __init__(...):
    # 其他代码不变
    self.relu1 = nn.ReLU(inplace=True)
    self.relu2 = nn.ReLU(inplace=True)
    self.relu3 = nn.ReLU(inplace=True)

修改后的forward函数：

def forward(self, x: Tensor) -> Tensor:
    identity = x

    out = self.conv1(x)
    out = self.bn1(out)
    out = self.relu1(out)

    out = self.conv2(out)
    out = self.bn2(out)
    out = self.relu2(out)

    out = self.conv3(out)
    out = self.bn3(out)

    if self.downsample is not None:
        identity = self.downsample(x)

    out += identity
    out = self.relu3(out)

    return out

这种方法优点是能用现有遍历方法直接获取所有层，缺点是需要修改模型源码，不适用于无法改动源码的场景。

内容的提问来源于stack exchange，提问作者gxxxh