如何合并两个PyTorch GNN模型类并解决参数传递错误？

错误原因分析 #

你的合并代码存在两个关键问题：

• forward方法仅定义了x和edge_index两个输入参数，但图分类任务需要传入batch_idx以区分不同图的节点，调用时传入3个参数（加上实例自身共4个），导致参数不匹配报错。
• 错误保留了原GraphGNNModel中self.GNN(x, edge_index)的调用逻辑，合并后已无self.GNN子模块；同时原GNN层的最终输出维度设置错误，应输出隐藏层维度c_hidden，再由头部层映射到分类任务的输出维度c_out。

正确合并后的代码 #

class GNNModel(nn.Module):
    
    def __init__(self, c_in, c_hidden, c_out, num_layers, activation_function, optimizer_name, learning_rate, dp_rate_linear, layer_name="GCN", **kwargs):
        """
        Inputs:
            c_in - 输入特征维度
            c_hidden - 隐藏层特征维度
            c_out - 输出特征维度（通常为分类任务的类别数）
            num_layers - 隐藏图卷积层的数量
            layer_name - 使用的图卷积层类型（GCN/GAT/GraphConv）
            dp_rate_linear - 网络中使用的Dropout概率
            kwargs - 图卷积层的额外参数（如GAT的注意力头数）
            activation_function - 激活函数
        """
        super().__init__()
        gnn_layer = gnn_layer_by_name[layer_name]
        
        layers = []
        # 兼容字符串或直接传入激活函数类的方式
        activation = eval(activation_function) if isinstance(activation_function, str) else activation_function
        
        in_channels = c_in
        # 构建隐藏层的图卷积+激活+Dropout序列
        for _ in range(num_layers - 1):
            layers.extend([
                gnn_layer(in_channels=in_channels, out_channels=c_hidden, **kwargs),
                activation(),
                nn.Dropout(p=dp_rate_linear)
            ])
            in_channels = c_hidden
        # 最后一层图卷积输出隐藏层维度，留给头部层做最终分类映射
        layers.append(gnn_layer(in_channels=in_channels, out_channels=c_hidden, **kwargs))
        self.layers = nn.ModuleList(layers)
        
        # 图分类头部：全局池化后执行线性分类
        self.head = nn.Sequential(
            nn.Dropout(p=dp_rate_linear),
            nn.Linear(c_hidden, c_out)
        )

    def forward(self, x, edge_index, batch_idx):
        # 执行图卷积层前向传播
        for l in self.layers:
            if isinstance(l, geom_nn.MessagePassing):
                x = l(x, edge_index)
            else:
                x = l(x)
        # 全局均值池化，将节点特征聚合为图级特征
        x = geom_nn.global_mean_pool(x, batch_idx)
        # 头部层输出最终结果
        x = self.head(x)
        return x

关键修改点说明 #

1. 参数修正：forward方法新增batch_idx参数，匹配图分类任务的输入需求。
2. 维度调整：最后一层图卷积的输出维度改为c_hidden，而非直接输出c_out，遵循先提取节点隐藏特征、再聚合为图特征分类的逻辑。
3. 逻辑修正：移除错误的self.GNN调用，直接通过self.layers完成图卷积层的前向传播，再依次执行全局池化和头部层计算。
4. 兼容性优化：对激活函数的处理增加了灵活性，支持直接传入函数类或字符串形式。

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