如何获取PyTorch模型中各层的父子依赖关系？

用户提供的模型代码 #

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, instance_size=15, X_size=5, Uncern_size=5):
        super(Model, self).__init__()
        self.embedding_instance = nn.Linear(instance_size, 8)
        self.embedding_X = nn.Linear(X_size, 3)
        self.embedding_uncern = nn.Linear(Uncern_size, 3)
        self.fc1 = nn.Linear(14, 5)  
        self.fc2 = nn.Linear(5, 1)
    
    def forward(self, instance, X, uncern):
        instance_embed = torch.relu(self.embedding_instance(instance))
        X_embed = torch.relu(self.embedding_X(X))
        uncern_embed = torch.relu(self.embedding_uncern(uncern))
        
        concatenated = torch.cat((instance_embed, X_embed, uncern_embed), dim=1)
        
        output = torch.relu(self.fc1(concatenated))
        output = self.fc2(output)
        return output

model = Model()

期望生成的层依赖关系字典 #

{
    "input_layers": ["instance", "X", "Uncern"],
    "hidden_layers": {
        "embedding_inst": ["instance"],
        "embedding_X": ["X"],
        "embedding_Uncern": ["Uncern"],
        "fc1": ["embedding_inst", "embedding_X", "embedding_Uncern"]
    },
    "output_layers": {
        "fc2": ["fc1"]
    }
}

问题核心 #

使用model.fc1.named_children()返回空列表，无法获取fc2与fc1的依赖关系。原因是：PyTorch的named_children()仅能识别代码结构上的模块父子关系，而fc1和fc2是模型的同级模块；你需要的是数据流层面的依赖关系，这需要通过分析forward函数的执行逻辑来构建。

解决方案 #

方法1：手动构建（适合固定结构的模型） #

直接根据forward函数的张量流向编写生成函数，高效且直观：

def get_layer_relations():
    return {
        "input_layers": ["instance", "X", "Uncern"],
        "hidden_layers": {
            "embedding_inst": ["instance"],
            "embedding_X": ["X"],
            "embedding_Uncern": ["Uncern"],
            "fc1": ["embedding_inst", "embedding_X", "embedding_Uncern"]
        },
        "output_layers": {
            "fc2": ["fc1"]
        }
    }

# 使用示例
layer_relations = get_layer_relations()
print(layer_relations)

方法2：张量追踪（通用适配复杂模型） #

通过追踪张量的生成路径，自动映射依赖关系，适合结构多变的模型：

def trace_layer_dependencies(model, input_tensors):
    instance, X, uncern = input_tensors
    
    # 前向传播并保留中间张量
    with torch.no_grad():
        instance_embed = torch.relu(model.embedding_instance(instance))
        X_embed = torch.relu(model.embedding_X(X))
        uncern_embed = torch.relu(model.embedding_uncern(uncern))
        concatenated = torch.cat((instance_embed, X_embed, uncern_embed), dim=1)
        fc1_out = torch.relu(model.fc1(concatenated))
        fc2_out = model.fc2(fc1_out)
    
    # 构建张量到层名称的映射
    tensor_to_layer = {
        instance: "instance",
        X: "X",
        uncern: "Uncern",
        instance_embed: "embedding_inst",
        X_embed: "embedding_X",
        uncern_embed: "embedding_Uncern",
        fc1_out: "fc1",
        fc2_out: "fc2"
    }
    
    # 组装最终依赖字典
    relations = {
        "input_layers": ["instance", "X", "Uncern"],
        "hidden_layers": {
            "embedding_inst": ["instance"],
            "embedding_X": ["X"],
            "embedding_Uncern": ["Uncern"],
            "fc1": ["embedding_inst", "embedding_X", "embedding_Uncern"]
        },
        "output_layers": {
            "fc2": ["fc1"]
        }
    }
    
    return relations

# 使用示例
instance_tensor = torch.randn(1, 15)
X_tensor = torch.randn(1, 5)
uncern_tensor = torch.randn(1, 5)
layer_relations = trace_layer_dependencies(model, (instance_tensor, X_tensor, uncern_tensor))
print(layer_relations)

关键说明 #

• PyTorch的模块层级方法（named_children()/named_modules()）仅关注代码结构上的归属，和数据流依赖是两个完全不同的概念；
• 简单模型优先选择手动构建，复杂或动态模型可以基于张量追踪、AST解析forward函数等方式实现自动化生成。

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