PyTorch中如何同步训练GNN内部实例化的线性神经网络

你当前给出的代码存在核心问题：将linear_NN的实例化逻辑写在了GNN的forward方法中，这会导致每次前向传播都会生成全新的linear_NN实例，参数无法持久化，自然也没法完成训练。正确的实现和训练流程如下：

第一步：修正模型类定义 #

要把内部子模型的实例化放到__init__方法中，这样PyTorch会自动把子模型的参数注册到父模型的参数列表中，不需要额外手动管理。

import torch
import torch.nn as nn
import torch.optim as optim

class linear_NN(nn.Module):
    def __init__(self, input_dim, out_dim):
        super().__init__()
        # 以下为线性网络的示例实现，可替换为你自己的逻辑
        self.fc1 = nn.Linear(input_dim, 128)
        self.act = nn.ReLU()
        self.fc2 = nn.Linear(128, out_dim)

    def forward(self, x, dim = 0):
        '''Forward pass'''
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return x

class GNN(nn.Module):
    def __init__(self, input_dim, n_hidden, out_dim, linear_in_dim, linear_out_dim):
        super().__init__()
        # GNN自身的层，以下为简化示例，可替换为你自己的GNN实现
        self.gnn_conv1 = nn.Linear(input_dim, n_hidden)
        self.gnn_conv2 = nn.Linear(n_hidden, out_dim)
        # 把linear_NN实例化放在__init__里，作为GNN的成员属性
        self.linear_submodel = linear_NN(linear_in_dim, linear_out_dim)

    def forward(self, h, dim = 0):
        '''Forward pass'''
        # 先跑GNN自身的前向逻辑
        h = self.gnn_conv1(h)
        h = nn.ReLU()(h)
        h = self.gnn_conv2(h)
        # 调用内部的线性子模型
        h = self.linear_submodel(h, dim)
        return h

第二步：训练流程 #

方案1：联合优化两个模型的所有参数（最常用） #

因为子模型参数已经自动注册到GNN的参数列表里了，直接把GNN实例的所有参数传给优化器即可，训练逻辑和普通单模型完全一致：

# 模型实例化
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = GNN(input_dim=16, n_hidden=32, out_dim=64, linear_in_dim=64, linear_out_dim=10).to(device)

# 优化器直接传入model.parameters()，会自动包含GNN和内部linear_NN的所有参数
optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

# 正常训练循环
for epoch in range(100):
    model.train()
    # 这里替换成你自己的数据集加载逻辑
    x = torch.randn(32, 16).to(device) # batch_size=32，特征维度16
    y = torch.randint(0, 10, (32,)).to(device)
    
    optimizer.zero_grad()
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()
    optimizer.step()
    
    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

方案2：分开优化两个模型的参数（适配差异化训练策略场景） #

如果需要给GNN和线性子模型设置不同的学习率、权重衰减等训练策略，可以分别提取两类参数传给不同的优化器：

model = GNN(input_dim=16, n_hidden=32, out_dim=64, linear_in_dim=64, linear_out_dim=10).to(device)

# 分别提取两类参数
gnn_params = list(model.gnn_conv1.parameters()) + list(model.gnn_conv2.parameters())
linear_params = model.linear_submodel.parameters()

# 两个优化器分别设置不同的超参数
optimizer_gnn = optim.Adam(gnn_params, lr=1e-3)
optimizer_linear = optim.Adam(linear_params, lr=5e-4)
loss_fn = nn.CrossEntropyLoss()

# 训练循环
for epoch in range(100):
    model.train()
    x = torch.randn(32, 16).to(device)
    y = torch.randint(0, 10, (32,)).to(device)
    
    optimizer_gnn.zero_grad()
    optimizer_linear.zero_grad()
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()
    # 两个优化器分别更新参数
    optimizer_gnn.step()
    optimizer_linear.step()
    
    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

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