在 PyTorch 中高效计算 PINN 求解器的导数

问题描述 投票：0回答：1

我目前正在尝试实现物理信息神经网络（PINNs）。 PINN 涉及计算模型输出相对于其输入的导数。然后使用这些导数来计算 PDE 残差，可以是 Heat、Burger、Navier-Stokes 方程等。因此，需要计算更高阶的偏导数。我尝试使用 

torch.autograd.grad
来计算这些偏导数。这是我实现的功能：

def diff(y, xs):
    grad = y
    ones = torch.ones_like(y)
    for x in xs:
        grad = torch.autograd.grad(grad, x, grad_outputs=ones, create_graph=True)[0]
    return grad


diff(y, xs)
只是计算 
y
相对于 
xs
中每个元素的导数。这种表示和计算偏导数的方法要容易得多：

y_pred = model([t, x, y])
diff(y_pred, [x, x]) # dydxx
diff(y_pred, [x, y, t]) # dydxyt

这个函数允许我以更紧凑和更简单的方式编写偏微分方程，这里的例子是一维热方程：

u = model([x, t])
def heat(u, x, t):
    alpha = 1.0
    return diff(u, [x, x]) - alpha * diff(u, [t])

PINN 除了数据之外的限制之一是训练时间。为了减少训练时间，我尝试使用 

torch.jit.script
、
torch.jit.trace
和 
torch.compile
。

  • JIT 脚本实现

    diff(y, xs)
    :

    @torch.jit.script
def D(y: torch.Tensor, xs: List[torch.Tensor]) -> torch.Tensor:
    func: torch.Tensor = y
    ones: List[Optional[torch.Tensor]] = [torch.ones_like(y)]
    for x in xs:
        grad = torch.autograd.grad(
            [func],
            [x],
            grad_outputs=ones,
            allow_unused=True,
            create_graph=True)[0]
        func = grad if grad is not None else torch.zeros_like(x)
    return func

  • 然后我也尝试实现热方程，但遇到了错误：

    • JIT 跟踪：

      N = 1000
x = torch.linspace(0.0, 1.0, N, requires_grad=True, device=device).reshape((N,1))
t = torch.linspace(0.0, 1.0, N, requires_grad=True, device=device).reshape((N,1))
u = torch.sin(x * np.pi) * torch.exp(-t**2) # A typical solution to the Heat Eqn.

def heat(u, x, t):
   alpha = 1.0
   res = D(u, [x, x]) - alpha * D(u, [t]) # Uses jit.script diff
   return res

traced_heat = torch.jit.trace(heat, (u, x, t))

      错误：

      Runtime  Error: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
   File "/tmp/ipykernel_19890/2258731968.py", line 6, in D
     ones: List[Optional[torch.Tensor]] = [torch.ones_like(y)]
     for x in xs:
         grad = torch.autograd.grad(
                ~~~~~~~~~~~~~~~~~~~ <--- HERE
             [func],
             [x],
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

我不知道是什么原因造成的，我没有使用 PyTorch 进行实验，但想导出我在 我的学习期间编写的内容（在 Tensorflow 中），因为 PyTorch 提供了更多的灵活性和控制。然而，有了如此大的灵活性，我不知道应该使用哪个选项。我应该尝试

TorchDynamo
吗？

  • 如果我在 

    torch._dynamo.explain
    函数上运行 
    heat
    ，我会得到以下解释：

    Dynamo produced 1 graphs with 0 graph break and 2 ops
Break reasons: 

TorchDynamo compilation metrics:
Function, Runtimes (s)
_compile, 0.0880, 0.0152, 0.0046
OutputGraph.call_user_compiler, 0.0000

  • 然后，如果我尝试编译 

    heat
    函数，我会收到以下错误：

    compheat = torch.compile(heat, backend='inductor', fullgraph=True)
r = compheat(u, x, t)

    错误：

    Unsupported: inlining disallowed: <function grad at 0x7f36642888b0>

from user code:
   File "/tmp/ipykernel_19890/543225163.py", line 3, in cheat
    res = D(u, [x, x]) - alpha * D(u, [t])
  File "/tmp/ipykernel_19890/2258731968.py", line 6, in D
    grad = torch.autograd.grad(


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True

我愿意接受设计建议和我所遇到的错误的解决方案。我是 PyTorch 的新手，对该框架的经验有限，但我渴望学习并应用其功能来解决复杂的问题，例如使用 PINN 的偏微分方程。由于我使用 PyTorch 的经验有限，如果我提供的任何信息不正确或不完整，我提前表示歉意。我来这里是为了学习，并感谢社区的任何更正或反馈。谢谢您的理解。

请注意，作为一个玩具示例，我使用了以下神经网络，它首先将输入作为列表，并在将其馈送到实际层之前将它们堆叠起来：

class NeuralNet(nn.Module):
    def __init__(self, n_input):
        super().__init__()
        self.linear_sigmoid_stack = nn.Sequential(
            nn.Linear(n_input, 32),
            nn.Sigmoid(),
            nn.Linear(32, 32),
            nn.Sigmoid(),
            nn.Linear(32, 1)
        )
    
    def forward(self, x):
        x = torch.hstack(x)
        return self.linear_sigmoid_stack(x)
python deep-learning pytorch neural-network jit
1个回答
0
投票

解决方案实际上非常简单，为了让 

torch.autograd
跟踪梯度，
torch.nn.Module.forward
必须接受 NN 输入作为单独的参数。这样就生成了后向图，并且可以使用 
torch.autograd
或 
torch.func.grad
来计算输出的梯度。

但是，请注意，

torch.func.grad
也需要模型参数作为模型的输入，因为训练需要梯度。参数也是如此。要使用 
torch.func.grad
计算问题中 
Heat
模型的梯度，

from functools import wraps
from torch import func
from torch import _dynamo as dynamo
model = Heat()

def make_sum(f):
    # Make functional summation
    @wraps(f)
    def _sum(*args, **kwargs):
        return f(*args, **kwargs).sum()

    return _sum

def fgrad(f, argnum=1):
    # Make functional gradient
    return func.grad(make_sum(f), argnums=argnum)

u = dynamo.allow_in_graph(model)
u_t = fgrad(u, argnum=1)
u_xx = fgrad(u_x, argnum=2)

def pde_residual(params, t, x):
    # Heat Equation gradient
    return u_t(params, t, x) - 0.05*u_xx(params, t, x)


pde_residual
可以使用
torch.compile
进行编译。为了对编译的训练进行基准测试，我还实现了 
torch.autograd
版本，

from torch import autograd

def grad(u, xs, create_graph=True, retain_graph=True):
    return autograd.grad(u.sum(),
                         xs,
                         create_graph=create_graph,
                         retain_graph=retain_graph)

def pde_residual(u_pred, t, x):
    u_t, u_x = func.grad(u_pred, [t, x])
    u_xx = func.grad(u_x, x)[0]
    return u_t - 0.05*u_xx

有趣的是，没有加速，事实上 

torch.autograd
的执行速度比编译版本快得多。

因此，我切换到动态图而不是编译并实现了一个名为 remin 的 PINN Python 包。以下是我的框架的功能。

remin 是一个 Python 模块，提供使用 PyTorch 开发物理信息神经网络 (PINN) 的框架。该模块集成了用于模型创建和训练的 PyTorch、用于几何创建的 Latin-Hypercube 采样，还包括一个用于创建几何的小型、用户友好的模块。

如果您想为 remin 做出贡献，请随时提交拉取请求或在 GitHub 上提出问题。

最新问题
© www.soinside.com 2019 - 2024. All rights reserved.