IT博客汇 | 「PINNs」求解Burgers方程

「PINNs」求解Burgers方程

Yacan Man发表于 2024-11-24 21:38:00

🖋️ PyTorch 的自动微分机制

在进行PINNs代码实战前，有必要对 PyTorch 的自动微分机制，也就是autograd 模块的实现进行复习。它使得神经网络中的梯度计算变得非常简单和高效。以下是复习自动微分机制的基本概念和步骤：

张量及梯度追踪 (`requires_grad=True`)

在 PyTorch 中，张量（Tensor）是主要的数据结构，类似于多维数组。自动微分机制会记录对张量的操作以构建计算图，但前提是张量需要启用梯度追踪。这可以通过设置 requires_grad=True 来完成：

x = torch.tensor(2.0, requires_grad=True)
t = torch.tensor(3.0, requires_grad=True)

这段代码创建了两个张量 x 和 t，并开启了梯度追踪。这样在计算过程中，PyTorch 会自动记录对 x 和 t 的操作，以便在之后计算梯度。

定义函数并构建计算图

PyTorch 会基于操作来构建计算图。每个操作都创建一个节点，并将输出作为另一个节点的输入。这样，计算图会记录下从输入到输出的所有依赖关系：

f = x ** 2 + 3 * x * t + t ** 2

这定义了一个函数：

$$ f(x, t) = x^2 + 3xt + t^2 $$

由于 x 和 t 都开启了梯度追踪，f 也会有梯度属性，并且 PyTorch 会自动构建计算图。

计算梯度

为了得到 f 对 x 和 t 的偏导数，我们使用 torch.autograd.grad 函数。它的常用参数包括：

outputs: 需要求导的输出张量，这里是 f。
inputs: 需要求导的变量，这里是 [x, t]。
grad_outputs: 当 outputs 是标量时通常设为 1，表示链式法则中的外部导数。
create_graph: 是否保留计算图。通常在二阶或更高阶导数计算时需要设为 True。

Tips / 提示
grad_outputs 参数在 PyTorch 中是用于指定目标张量 outputs 的“外部导数”（或梯度）：在微积分中，链式法则用于计算复合函数的导数。假设我们有两个变量 $y$ 和 $z$ ，且它们之间存在如下依赖关系：
$$ z = f(y), \quad y = g(x) $$
那么，如果我们要计算$\frac{\partial z}{\partial x}$，可以使用链式法则得到：
$$ \frac{\partial z}{\partial x} = \frac{\partial z}{\partial y} \cdot \frac{\partial y}{\partial x} $$
在 PyTorch 的 autograd 中，如果我们想要计算 z 对 x 的梯度，grad_outputs 就提供了“外部导数” $\frac{\partial z}{\partial y}$，从而帮助 autograd 计算梯度。这里设置设置grad_outputs=torch.ones_like(f)其实就是另“外部导数” $\frac{\partial z}{\partial y}=1$.

gradients = torch.autograd.grad(
    outputs=f,
    inputs=[x, t],
    grad_outputs=torch.ones_like(f),
    create_graph=True
)

这行代码会返回 f 对 x 和 t 的偏导数，并将结果存储在 gradients 中。

获取偏导数结果

print("df/dx:", gradients[0])  # df/dx: tensor(13., grad_fn=<AddBackward0>)
print("df/dt:", gradients[1])  # df/dt: tensor(12., grad_fn=<AddBackward0>)

这将显示 f 对 x 和 t 的偏导数。根据上面 $f(x, t)$ 的定义，当 $x=2, t=3$ 时，可得：

$$ \frac{\partial f}{\partial x} = 2x+3t=13, \frac{\partial f}{\partial t} = 3x+2t=12 $$

矩阵、向量求导

在 PyTorch 中，对矩阵或张量求导的过程可以理解为对每个元素进行逐元素（element-wise）求导。

假设我们有两个张量 X 和 u，其中 u 是 X 的函数。X 的形状是 (2, 2)，即：

# 假设 x 和 t 是这些值，直接设置 requires_grad=True
x = torch.tensor([2.0, 3.0], requires_grad=True)
t = torch.tensor([0.0, 1.0], requires_grad=True)


from torch.autograd import Variable

# 创建叶子节点副本
X = Variable(torch.stack([x, t], dim=1), requires_grad=True)

u 的形状与 X 相同，因为 u 是对 X 的逐元素操作：

u = X ** 2 + 2 * X

对于矩阵 u 中的每个元素 u[i, j]，它是 X[i, j] 的函数。因此，计算 u 对 X 的梯度相当于逐元素对 X 中的每个元素计算偏导数。这会生成一个与 X 相同形状的矩阵，每个元素对应于 u 的某个元素对 X 相应元素的偏导数。

torch.autograd.grad 可以对 u 中的每个元素求 X 中每个元素的偏导。由于 u 和 X 是逐元素关联的，torch.autograd.grad 会计算每个 u[i, j] 对应的 X[i, j] 的梯度。

在代码中，grad_outputs=torch.ones_like(u) 用来设置“外部导数”，这相当于在链式法则中使用单位向量来表示标量梯度。具体来说，对于 u[i, j] 对 X[i, j] 的偏导数，我们设定 grad_outputs[i, j] = 1，即每个元素都作为独立的标量。

du_dX = torch.autograd.grad(
    inputs=X,       # 计算 X 的梯度
    outputs=u,      # u 是基于 X 的函数
    grad_outputs=torch.ones_like(u),  # 设置外部导数为 1
    retain_graph=True,
    create_graph=True,
    allow_unused=True
)[0]

du_dX[:, 0] 和 du_dX[:, 1] 分别给出了 u 对 X 的两个列的偏导：

du_dX[:, 0] 表示 u 对 x 的偏导数。
du_dX[:, 1] 表示 u 对 t 的偏导数。

print("du/dx:", du_dX[:, 0])  # du/dx: tensor([6., 8.], grad_fn=<SelectBackward0>)
print("du/dt:", du_dX[:, 1])  # du/dt: tensor([2., 4.], grad_fn=<SelectBackward0>)

🔛 求解过程

Intro

Burgers方程是一类非线性偏微分方程，可以用于模拟流体动力学、非线性声学和交通流等。在论文中，PINNs使用自动微分计算网络输出对输入的导数，并将其替代到Burgers方程中形成损失函数。通过最小化损失函数，使得神经网络的输出逼近真实解。

RAISSI M, PERDIKARIS P, KARNIADAKIS G 2017. Physics Informed Deep Learning (Part I): Data-driven Solutions of Nonlinear Partial Differential Equations.

Burgers方程的具体形式为：

$$ u_t + u u_x - \left(\frac{0.01}{\pi}\right) u_{xx} = 0, x \in [-1, 1], \, t \in [0, 1] $$

生成空间-时间网格

h, k = .1, .1
x = torch.arange(-1, 1 + h, h)  # len(x) = 21, 空间维度
t = torch.arange(0, 1 + k, k)  # len(t) = 11, 时间维度

x 和 t 分别定义了空间和时间维度的点，通过步长 h 和 k 划分区间。这里 x 的范围为 [-1, 1]，而 t 的范围为 [0, 1]。

X = torch.stack(torch.meshgrid(x, t, indexing='ij')).reshape(2, -1).T

使用 torch.meshgrid 生成了一个二维的网格，它包含了空间和时间上所有组合的点。然后用 torch.stack 将 x 和 t 的网格组合成 [len(x) * len(t), 2] 的张量 X。该张量的每一行表示一个 (x, t) 组合，这样的结构便于后续计算。

定义边界条件

边界条件和初始条件的定义是PINNs的关键步骤。对于Burgers方程，训练数据包括空间上的边界条件和时间上的初始条件：

bc1 = torch.stack(torch.meshgrid(x[0], t, indexing='ij')).reshape(2, -1).T  # 下边界点
bc2 = torch.stack(torch.meshgrid(x[-1], t, indexing='ij')).reshape(2, -1).T  # 上边界点
ic = torch.stack(torch.meshgrid(x, t[0], indexing='ij')).reshape(2, -1).T  # 起始时刻点

bc1 和 bc2 分别代表空间维度上边界的 (x, t) 点，其中 bc1 对应 x=-1 处的下边界，bc2 对应 x=1 处的上边界。
ic 表示初始条件的采样点，即 t=0 时刻整个空间维度的所有点 (x, t=0)。

合并训练数据

将 bc1、bc2 和 ic 合并成 X_train，它包含了所有边界和初始条件的点，用于训练PINN模型使其满足这些物理边界和初始条件。

X_train = torch.cat([bc1, bc2, ic])  # X_train

定义目标值（真实解）

y_bc1 = torch.zeros(len(bc1))
y_bc2 = torch.zeros(len(bc2))
y_ic = -torch.sin(math.pi * ic[:, 0])
y_train = torch.cat([y_bc1, y_bc2, y_ic]).unsqueeze(1)

y_bc1 和 y_bc2 分别是上下边界的目标值，这里均为0，即满足边界条件 $u(t, -1) = u(t, 1) = 0$。
y_ic 是初始条件的目标值，由 $u(0, x) = -\sin(\pi x)$ 确定。
将 y_bc1、y_bc2 和 y_ic 合并成 y_train，以与 X_train 中的点对应，最终构成神经网络的目标输出。

数据迁移到GPU

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
X = X.to(device)
X.requires_grad = True
X_train = X_train.to(device)
y_train = y_train.to(device)

将数据迁移到GPU（如果可用），加速计算。
设置 X.requires_grad = True 以启用自动微分，因为Burgers方程中涉及到对网络输出 $u(t, x)$ 关于时间和空间的偏导数。

多层感知机

首先建立一个深度多层感知机用于近似非线性函数$u(t,x)$：

class MultiLayerPerceptron(nn.Module):
    """
    Multi-layer Perceptron.
    """

    def __init__(
            self,
            input_size,
            hidden_size,
            output_size,
            depth,
            act=torch.nn.Tanh,
    ):
        """
        :param input_size: input size.
        :param hidden_size: hidden size.
        :param output_size: output size.
        :param depth: depth of the network.
        :param act: activation function.
        """
        super(MultiLayerPerceptron, self).__init__()

        layers = [('Input', torch.nn.Linear(input_size, hidden_size)), ('Input_Activation', act())]
        for i in range(depth):
            layers.append(
                ('Hidden_%d' % i, torch.nn.Linear(hidden_size, hidden_size))
            )
            layers.append(('Activation_%d' % i, act()))
        layers.append(('Output', torch.nn.Linear(hidden_size, output_size)))

        self.layers = torch.nn.Sequential(OrderedDict(layers))

    def forward(self, x):
        out = self.layers(x)
        return out

该 MLP 的层数和激活函数类型都可以灵活调整，便于适应不同的应用需求。

建立PINNs网络

class PINNForBurgers:
    def __init__(self, device, X, X_train, y_train, model):

        self.model = model

        self.X = X
        self.X_train = X_train
        self.y_train = y_train

        self.criterion = torch.nn.MSELoss()
        self.iter = 1

        self.optimizer = torch.optim.LBFGS(
            self.model.parameters(),
            lr=1.0,
            max_iter=50000,
            max_eval=50000,
            history_size=50,
            tolerance_grad=1e-7,
            tolerance_change=1.0 * np.finfo(float).eps,
            line_search_fn="strong_wolfe",  # better numerical stability
        )

        self.adam = torch.optim.Adam(self.model.parameters())

    def loss_func(self):
        # this is more like a not so elegant hack to zero grad both optimizers
        self.adam.zero_grad()
        self.optimizer.zero_grad()

        y_pred = self.model(self.X_train)
        loss_data = self.criterion(y_pred, self.y_train)
        u = self.model(self.X)

        du_dX = torch.autograd.grad(
            inputs=self.X,  # The shape of X: [len(x) * len(t), 2], X的第一列是空间坐标x, 第二列是是时间坐标t.
            outputs=u,
            grad_outputs=torch.ones_like(u),
            retain_graph=True,
            create_graph=True
        )[0]

        du_dt = du_dX[:, 1]
        du_dx = du_dX[:, 0]
        du_dxx = torch.autograd.grad(
            inputs=self.X,
            outputs=du_dX,
            grad_outputs=torch.ones_like(du_dX),
            retain_graph=True,
            create_graph=True
        )[0][:, 0]

        loss_pde = self.criterion(du_dt + u.squeeze() * du_dx, 0.01 / math.pi * du_dxx)

        loss = loss_pde + loss_data
        loss.backward()
        if self.iter % 100 == 0:
            print(self.iter, loss.item())
        self.iter = self.iter + 1
        return loss

    def train(self):
        self.model.train()
        for i in range(1000):
            self.adam.step(self.loss_func)
        self.optimizer.step(self.loss_func)

    def eval_(self):
        self.model.eval()

这里面最主要的就是损失函数 (loss_func)的建立，这里采用的是[KUMAR V, GLEYZER L, KAHANA A, et al. MyCrunchGPT: A chatGPT assisted framework for scientific machine learning[C]//,202310.48550/arXiv.2306.15551.](https://www.researchgate.net/publication/371908937_CrunchGPT_A_chatGPT_assisted_framework_for_scientific_machine_learning)提到的方法：

也就是在损失函数中加入物理损失来计算求解。

这里的损失函数定义为 loss = loss_pde + loss_data。其中，loss_data也就是已知边界点的计算误差：self.criterion(y_pred, self.y_train)；loss_pde则是指物理损失 self.criterion(du_dt + u.squeeze() * du_dx, 0.01 / math.pi * du_dxx)，也就是方程 $u_t + u u_x - \left(\frac{0.01}{\pi}\right) u_{xx} = 0$ 的误差。

训练结果

从最后的结果可以看出，我们根据已知边界点（x符号）的数据信息以及物理方程的定义拟合出了非线性函数 $u(t,x)$ 的近似表达式。

提取不同时刻下对应的函数解：

🪛 方程常数项的学习

如上面所述，Burgers方程的具体形式为：

$$ u_t + u u_x - \left(\frac{0.01}{\pi}\right) u_{xx} = 0, x \in [-1, 1], \, t \in [0, 1] $$

但是很多工程问题我们是并不能确切地使用完整的物理方程来进行描述的。例如，如果我们不确定 $u u_x$ 与 $u_{xx}$ 前的系数，那么我们也可以借助于神经网络强大的学习能力来进行参数“拟合”。

class PhysicsInformedNN:
    """
    The physics-guided neural network.
    """

    def __init__(self, X, u, model, device):

        # Data
        self.x, self.t = self.handle_m(X)
        self.u = torch.tensor(u).float().to(device)

        # PED params
        self.lambda_1, self.lambda_2 = [torch.nn.Parameter(
            torch.tensor(
                [v],
                requires_grad=True
            ).to(device)
        ) for v in [0.0, -6.0]]

        self.nn = model
        # 是 lambda_1 与 lambda_2 成为模型的一部分, 从而被优化器自动管理, 并在训练过程中被更新.
        for name, param in {'lambda_1': self.lambda_1, 'lambda_2': self.lambda_2}.items():
            self.nn.register_parameter(name, param)
            
     # ...

Warning / 注意
在 Burgers 方程里，$u_{xx}$ 前面的常数项被称为扩散系数（diffusion coefficient），其控制了扩散或扩散速率的强弱，且永远是正值，这是因为如果扩散系数是负值，会导致“负扩散”现象，即不均匀性反而增加，这是不符合物理实际的。那我们应该如何限制自定义参数永远为正值呢？
f = u_t + self.lambda_1 * u * u_x - torch.exp(self.lambda_2) * u_xx
其实，只要使用torch.exp()来将其包裹起来，根据指数函数的特点，无论self.lambda_2在训练过程中如何变化，$u_xx$ 前的系数torch.exp(self.lambda_2)永远都为正值。

具体的实现方式就是将 $u u_x$ 与 $u_{xx}$ 的系数设置为可学习参数，将其绑定注册到神经网络中，即可在神经网络的训练过程中进行参数的“拟合”。

最后的结果对比证明，当原始数据不存在噪音的情况下，PINN基本可以“完美地”拟合 $u u_x$ 与 $u_{xx}$ 的常数项，其真实值分别为 1 和 $\frac{0.01}{\pi}$ ，拟合的结果为0.99952和0.0031969，当数据存在1%的噪音时，拟合精度显著下降。

💡 一个小Trick

在PINNs-Torch: Enhancing Speed and Usability of Physics-Informed Neural Networks with PyTorch提供的代码中，作者使用了个小Trick——采用两个优化器（LBFGS 和 Adam）对训练过程中的PINN进行优化。定义如下：

self.optimizer = torch.optim.LBFGS(
    self.nn.parameters(),
    lr=1.0,
    max_iter=50000,
    max_eval=50000,
    history_size=50,
    tolerance_grad=1e-5,
    tolerance_change=1.0 * np.finfo(float).eps,
    line_search_fn="strong_wolfe"  # can be "strong_wolfe"
)

self.optimizer_Adam = torch.optim.Adam(self.nn.parameters())

在训练函数中的具体调用方式为：

def train(self, epoch_num: int):
    self.nn.train()
    for n_iter in range(epoch_num):
        # ...
        loss = ...

        # Backward and optimize
        self.optimizer_Adam.zero_grad()
        loss.backward()
        self.optimizer_Adam.step()

        # ...

    # Backward and optimize
    self.optimizer.step(self.loss_func)

使用两个优化器（LBFGS 和 Adam）是设计PINN网络常见的优化策略，目的是结合两种优化算法的优点，达到快速收敛和高精度优化的效果。

两种优化器的特点：

Adam 优化器
- 优点
  - 快速收敛：Adam 是一种自适应优化算法，能够根据梯度的历史记录动态调整学习率，在优化初期表现优异。
  - 鲁棒性：适合处理非凸问题和梯度震荡，尤其是参数较多或损失函数复杂时。
  - 全局探索：Adam 能够快速探索参数空间，找到一个接近最优解的区域。
- 不足:
  - 在优化后期，Adam 的更新步长可能过大，难以细致地调整参数，导致优化精度不足。
LBFGS 优化器
- 优点:
  - 高精度：LBFGS 是一种拟牛顿方法，利用二阶信息（Hessian 矩阵的近似）进行优化，收敛速度快且优化结果精确。
  - 适合平滑问题：在损失函数较平滑时（例如 PINNs 中的物理残差），LBFGS 可以很好地找到局部最优解。
- 不足:
  - 计算开销大：由于涉及近似二阶导数的计算，LBFGS 的开销较大。
  - 对初始点敏感：如果初始点较差，LBFGS 可能陷入次优解。

使用两个优化器是为了取长补短，结合两者的优势。

🧰 半监督神经网络

物理信息神经网络本质上是一种半监督学习方法，因为它结合了监督学习和无监督学习的特性：

监督学习部分：使用已知的边界条件和初始条件数据进行监督学习，self.x_u 和 self.t_u: 边界条件点的位置，self.u: 这些点上已知的精确解。
无监督学习部分：PINNs 在定义域内随机采样伪点这些点没有已知的解，但它们需要满足偏微分方程的物理约束，计算Loss的过程中，通过让物理方程的误差趋近于0来进行梯度优化。

例如上图中，我们可以利用边界点（已知对应 $u$ 值的点）来进行计算神经网络部分的Loss，然后在平面内随机采样一些点（不知道对应 $u$ 值，但这些点预测的 $u$ 要符合物理方程）作为物理信息部分的Loss。