如何在PyTorch中计算偏导数？多层神经网络下对输入变量的一阶及二阶偏导实现问询

嘿，这个需求在物理信息神经网络（PINN）场景里太常见了！我来一步步带你实现对u(x,t)求关于x的一阶偏导、t的一阶偏导，以及x的二阶偏导的操作。

第一步：定义你的MLP模型 #

首先先把你描述的2×50、50×50、50×1的网络结构用PyTorch实现出来，这里假设你用ReLU激活（你可以换成自己需要的激活函数）：

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(2, 50),
            nn.ReLU(),
            nn.Linear(50, 50),
            nn.ReLU(),
            nn.Linear(50, 1)
        )
    
    def forward(self, x_t):
        # x_t是形状为[batch_size, 2]的张量，第一列是x，第二列是t
        return self.network(x_t)

第二步：准备输入并前向传播 #

要注意输入张量必须开启梯度追踪（requires_grad=True），这样PyTorch才能记录计算图用于求导：

# 初始化模型
model = MLP()

# 准备输入：x和t各是[100,1]的张量，拼接成[100,2]的输入
x = torch.randn(100, 1, requires_grad=True)
t = torch.randn(100, 1, requires_grad=True)
x_t = torch.cat([x, t], dim=1)  # 形状：[100, 2]

# 前向传播得到u
u = model(x_t)  # 形状：[100, 1]

第三步：计算一阶偏导（∂u/∂x 和 ∂u/∂t） #

我们用torch.autograd.grad()来计算偏导，这个方法可以直接返回目标输出对指定输入的梯度，而且支持批量计算：

# 计算一阶偏导
du_dx, du_dt = torch.autograd.grad(
    outputs=u,                # 要对其求导的输出张量
    inputs=[x, t],            # 求导的目标输入
    grad_outputs=torch.ones_like(u),  # 对应每个输出元素的权重（这里取1表示直接求导）
    create_graph=True,        # 保留计算图，方便后续计算二阶导
    retain_graph=True         # 可选：如果之后还要复用u的计算图则保留
)

# 此时du_dx和du_dt的形状都是[100,1]，对应每个样本的一阶偏导
print(f"du/dx shape: {du_dx.shape}")
print(f"du/dt shape: {du_dt.shape}")

第四步：计算x的二阶偏导（∂²u/∂x²） #

因为我们之前设置了create_graph=True，所以du_dx的计算图也被保留了，可以直接对它再次求导：

# 计算二阶偏导d²u/dx²
d2u_dx2 = torch.autograd.grad(
    outputs=du_dx,
    inputs=x,
    grad_outputs=torch.ones_like(du_dx),
    create_graph=False  # 如果不需要更高阶导，这里可以设为False节省内存
)[0]  # autograd.grad返回的是元组，取第一个元素就是梯度张量

print(f"d²u/dx² shape: {d2u_dx2.shape}")  # 形状：[100,1]

关键注意事项 #

• 梯度追踪：所有需要求导的输入张量必须设置requires_grad=True，否则PyTorch不会记录它们的计算图。
• create_graph参数：如果需要计算高阶导数，必须在求一阶导时设置create_graph=True，这样一阶导数的计算图才会被保留。
• 批量处理：上述代码支持批量输入（这里是100个样本），每个样本的偏导都会被独立计算，输出形状和输入一致。
• 激活函数影响：如果你用了非光滑的激活函数（比如ReLU），在不可导点处的梯度会是0，这是正常现象，如果需要更平滑的梯度可以换成Swish、GELU等激活函数。

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