PyTorch物理信息神经网络全批次输出单一值问题排查

问题概述 #

物理信息神经网络（PINN）以距离x和时间t为输入，输出压力h和流量q，训练时将管道流体控制PDE的求解结果作为loss_pde纳入损失函数。但输入变化时（x固定为155m，t为0-10s的380个值），网络全批次始终输出单一值，数据损失仅能降至约163，无法拟合预期的时序变化曲线。已尝试调整激活函数、优化器及学习率，问题未解决。

现有代码细节 #

网络架构 #

class FCN(nn.Module):
    def __init__(self, N_INPUT, N_OUTPUT, N_HIDDEN, N_LAYERS):
        super().__init__()
        activation = nn.Tanh
        self.fcs = nn.Sequential(
            nn.Linear(N_INPUT, N_HIDDEN),
            activation()
        )
        self.fch = nn.Sequential(*[
            nn.Sequential(
                nn.Linear(N_HIDDEN, N_HIDDEN),
                activation()
            ) for _ in range(N_LAYERS - 1)
        ])
        self.fce = nn.Linear(N_HIDDEN, N_OUTPUT)

    def forward(self, x, t):
        inputs = torch.cat([x, t], axis = 1)
        inputs = self.fcs(inputs)
        inputs = self.fch(inputs)
        outputs = self.fce(inputs)
        return outputs

pinn = FCN(2, 2, 9, 12)

原始训练循环 #

for i in range (200001):

    w_f = 1e-3

    optimiser.zero_grad()

    colloc_output = pinn(x_colloc, t_colloc)

    h_hat, q_hat = colloc_output[:,0], colloc_output[:,1]

    dq_dt = torch.autograd.grad(q_hat, t_colloc, torch.ones_like(q_hat), create_graph=True)[0]
    dq_dx = torch.autograd.grad(q_hat, x_colloc, torch.ones_like(q_hat), create_graph=True)[0]
    dh_dt = torch.autograd.grad(h_hat, t_colloc, torch.ones_like(h_hat), create_graph=True)[0]
    dh_dx = torch.autograd.grad(h_hat, x_colloc, torch.ones_like(h_hat), create_graph=True)[0]

    F1 = Cs_A * dq_dt + q_hat * dq_dx + g * Cs_A**2 * dh_dx + f * (torch.abs(q_hat) * q_hat) / (2 * diam)
    F2 = Cs_A * dh_dt + q_hat * dh_dx + a**2/g * dq_dx

    loss_pde = torch.mean(F1**2 + F2**2)

    train_output = pinn(x_train_rand, t_train_rand)

    loss_data = torch.mean((train_output - hq_tr)**2)

    loss = w_f * loss_pde + loss_data 

    loss.backward()
    optimiser.step()

更新后的Dataset与DataLoader #

class TrainData(Dataset):
    def __init__(self, x_traindata, t_traindata, h_train, q_train):
        train_input = torch.cat([x_traindata, t_traindata], axis = 1)
        train_aim = torch.cat([h_train, q_train], axis = 1)
        self.train_input = train_input
        self.train_aim = train_aim
        
    def __len__(self):
        return self.train_input.size(0)

    def __getitem__(self, idx):
        input_value_x = self.train_input[idx, 0]
        input_value_t = self.train_input[idx, 1] 
        output_value_h = self.train_aim[idx, 0]
        output_value_q = self.train_aim[idx, 1]
        return input_value_x, input_value_t, output_value_h, output_value_q

class CollocData(Dataset):
    def __init__(self, x_colloc, t_colloc):
        colloc_input = torch.cat([x_colloc, t_colloc], axis = 1)
        self.colloc_input = colloc_input
        
    def __len__(self):
        return self.colloc_input.size(0)
    
    def __getitem__(self, idx):
        input_x = self.colloc_input[idx, 0]
        input_t = self.colloc_input[idx, 1]
        return input_x, input_t

输入定义与DataLoader初始化 #

x_train = torch.tensor(x_inp.values).to(dtype=torch.float32)
t_train = torch.tensor(t_inp.values).to(dtype=torch.float32)
h_train = torch.tensor(h_inp.values).to(dtype=torch.float32)
q_train = torch.tensor(q_inp.values).to(dtype=torch.float32)

x_colloc = torch.linspace(0, 300, 500).view(-1,1).requires_grad_(True)
t_colloc = torch.linspace(0, 10, 500).view(-1,1).requires_grad_(True)

test_dataset = TrainData(x_train, t_train, h_train, q_train)
colloc_dataset = CollocData(x_colloc, t_colloc)

train_dataloader = DataLoader(test_dataset, batch_size = x_train.size(0), shuffle=True)
colloc_dataloader = DataLoader(colloc_dataset, batch_size = x_colloc.size(0), shuffle = True)

更新后的训练循环开头 #

for i in range (200001):
    for train_batch, colloc_batch in zip(train_dataloader, colloc_dataloader):

        optimiser.zero_grad()
        
        train_x, train_t, target_h, target_q = train_batch
        colloc_x, colloc_t = colloc_batch
            
        colloc_output = pinn(colloc_x.unsqueeze(1), colloc_t.unsqueeze(1))
            
        h_hat, q_hat = colloc_output[:,0], colloc_output[:,1]

排查方向与解决建议 #

1. 输入数据归一化/标准化 #

• 输入x固定为155，t为0-10，数值范围差异大，导致网络难以捕捉t的变化信息。
• 解决：对x和t进行标准化（减去均值除以标准差），确保输入特征处于同一数量级。

2. 损失函数权重失衡 #

• 当前w_f = 1e-3，PDE损失权重远低于数据损失，网络可能优先拟合数据均值而非时序变化。
• 解决：先将w_f调至1或更高，观察两类损失的变化，再逐步平衡权重比例。

3. 深层网络梯度消失 #

• 12层隐藏层搭配Tanh激活函数，易出现梯度消失，导致参数无法有效更新。
• 解决：
  • 替换激活函数为ReLU/GELU；
  • 在每个隐藏层后添加nn.LayerNorm(N_HIDDEN)；
  • 暂时减少隐藏层数量（如4-6层），验证是否因过深导致梯度问题。

4. 自动求导维度检查 #

• 若torch.autograd.grad返回的导数维度与输入不匹配，会导致PDE损失计算异常，网络无法学习变化规律。
• 解决：打印dq_dt、dh_dx等导数的形状，确保与q_hat、h_hat的形状一致。

5. 小批次训练替代全批次 #

• 全批次输入易让网络陷入输出均值的局部最优，缺乏训练随机性。
• 解决：改用小批次训练（如batch_size=32/64），帮助网络跳出局部最优。

6. 参数初始化优化 #

• 默认线性层初始化可能导致初始输出过于一致，难以学习变化。
• 解决：手动采用Xavier/Kaiming初始化：

  def init_weights(m):
      if isinstance(m, nn.Linear):
          torch.nn.init.xavier_uniform_(m.weight)
          m.bias.data.fill_(0.01)
  pinn.apply(init_weights)
  

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