如何查询PyTorch中层的输入节点/层名称 以torch.cat输入来源查询为例

PyTorch默认采用动态计算图机制，不会像Keras的静态图那样预存所有层的输入输出元信息，因此没有直接对应model.get_layer(id=idx).input.name的原生接口，你可以根据使用场景选择以下三种方案实现需求：

方案1：前向传播钩子（运行时查询，最常用） #

你可以通过PyTorch提供的register_forward_hook接口为指定层绑定回调函数，层完成前向计算时会自动触发钩子，捕获对应的输入、输出数据，结合层的唯一名称即可关联输入的来源。

实现步骤 #

1. 首先为模型所有层绑定唯一标识，注册钩子：

# 存储层的输入输出信息
layer_io_record = {}
# 模块到名称的映射
module_to_name = {}

def io_hook(module, input, output):
    layer_name = module_to_name[module]
    layer_io_record[layer_name] = {
        "inputs": input,
        "output": output
    }

# 遍历模型所有子层注册钩子
model = LitClassifier()
for name, module in model.named_modules():
    if name != "": # 跳过根模块
        module_to_name[module] = name
        module.register_forward_hook(io_hook)

2. 如果需要捕获torch.cat的输入，可以把torch.cat封装为自定义层加入模型，方便绑定钩子：

# 自定义Cat层
class CatLayer(nn.Module):
    def __init__(self, dim=1):
        super().__init__()
        self.dim = dim
    def forward(self, *tensors):
        return torch.cat(tensors, dim=self.dim)

# 修改原模型：在__init__中添加
self.cat_layer = CatLayer(dim=1)
# 在forward中替换原torch.cat调用
x = self.cat_layer(img, tab)

3. 运行一次前向传播触发钩子后，即可通过比对torch.cat层的输入和其他层的输出，匹配到来源层。

方案2：计算图回溯（无需改模型，适合临时排查） #

完成前向传播后，可以通过张量的.grad_fn属性回溯动态计算图，直接找到输入张量对应的生成节点：

# 运行前向传播得到输出
test_img = torch.randn(1, 3, 224, 224) # 替换为你的输入尺寸
test_tab = torch.randn(1, 5)
output = model(test_img, test_tab)

# 回溯到cat操作的节点
cat_grad_fn = output.grad_fn.next_functions[0][0] # 根据你的计算图层级调整索引
# 打印所有输入的来源操作
for idx, (input_fn, _) in enumerate(cat_grad_fn.next_functions):
    print(f"第{idx}个输入的来源操作：{input_fn.name()}")

该方案不需要修改模型代码，但底层节点命名相对抽象，需要自行映射到模型的自定义层名称，适合简单模型快速排查问题。

方案3：TorchFX静态图转换（结构化分析，无需运行前向） #

TorchFX是PyTorch官方提供的图转换工具，可以把动态图转换为可遍历的静态图结构，直接获取每层的输入输出依赖关系，是最接近Keras静态查询能力的方案：

import torch.fx

# 把模型转换为FX静态图
traced_model = torch.fx.symbolic_trace(model)
# 遍历所有图节点
for node in traced_model.graph.nodes:
    # 匹配torch.cat操作节点
    if node.op == "call_function" and node.target == torch.cat:
        print(f"cat操作的输入来源节点：{[n.name for n in node.all_input_nodes]}")

该方案不需要运行前向传播，即可结构化获取整个模型的层依赖关系，适合复杂模型的批量分析场景。

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