如何在多输出Physics Informed Neural Network中并行计算导数

我正在开发一个双输入、N输出的物理信息神经网络（PINN），损失函数需要用到输出关于输入x和y的高阶导数。现有代码仅支持单输出场景，现在需要在不使用for循环的前提下，并行计算每个输出列的导数，确保每个导数（如dudx、dudy等）的形状与输出u（[batch_size, N]）一致。

现有单输出代码如下：

def gradient(y, x, grad_outputs=None):
    if grad_outputs is None:
        grad_outputs = torch.ones_like(y)
    grad = torch.autograd.grad(y, [x], grad_outputs=grad_outputs, create_graph=True)[0]
    return grad


def compute_derivatives(x, y, u):
    dudx = gradient(u, x)
    dudy = gradient(u, y)

    dudxx = gradient(dudx, x)
    dudyy = gradient(dudy, y)

    dudxxx = gradient(dudxx, x)
    dudxxy = gradient(dudxx, y)
    dudyyy = gradient(dudy, y)

    dudxxxx = gradient(dudxxx, x)
    dudxxyy = gradient(dudxxy, y)
    dudyyyy = gradient(dudyyy, y)

    return dudxx, dudyy, dudxxxx, dudyyyy, dudxxyy

核心思路 #

PyTorch的torch.autograd.grad本身支持批量输出的导数计算，关键是让grad_outputs的形状与输出u匹配，同时确保求导过程中保留维度信息，避免自动降维。

修改后的实现代码 #

import torch

def gradient(y, x, grad_outputs=None):
    if grad_outputs is None:
        # 针对多输出生成与y形状完全一致的全1张量，保证每个输出维度的导数独立计算
        grad_outputs = torch.ones_like(y)
    # 保留计算图以支持高阶导数，retain_graph可选（按需调整）
    grad = torch.autograd.grad(y, x, grad_outputs=grad_outputs, create_graph=True, retain_graph=True)[0]
    # 确保导数形状与y一致：若输入x是单维度张量，扩展维度匹配u的[N]输出列
    if grad.dim() < y.dim():
        grad = grad.expand_as(y)
    return grad


def compute_derivatives(x, y, u):
    # 一阶导数：形状[batch_size, N]
    dudx = gradient(u, x)
    dudy = gradient(u, y)

    # 二阶导数
    dudxx = gradient(dudx, x)
    dudyy = gradient(dudy, y)

    # 三阶导数（修正原代码的错误：dudyyy需基于dudyy求导）
    dudxxx = gradient(dudxx, x)
    dudxxy = gradient(dudxx, y)
    dudyyy = gradient(dudyy, y)

    # 四阶导数
    dudxxxx = gradient(dudxxx, x)
    dudxxyy = gradient(dudxxy, y)
    dudyyyy = gradient(dudyyy, y)

    return dudxx, dudyy, dudxxxx, dudyyyy, dudxxyy

关键要点说明 #

• 并行计算支持：torch.ones_like(y)让PyTorch自动并行处理N个输出列的导数，无需手动循环。
• 维度一致性：通过expand_as(y)将单维度导数结果扩展为[batch_size, N]，避免因输入x/y的单维度特性导致导数降维。
• 高阶导数连续性：create_graph=True保留计算图，确保后续高阶导数能正常计算；retain_graph=True可选，若无需重复使用当前计算图可移除。
• 原代码修正：修正了原代码中dudyyy的求导对象错误，保证导数阶数的正确性。

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