【Pytorch: Deepwave】GPU引领的正演与数值全波形反演 (FWI)

0. Deepwave的基本介绍

Deepwave是一个强大的使用有限差分法实现波动方程的正演和数值迭代反演的Python库.
相比于一般的基于Matlab或者常规代码实现的正演或者数值迭代反演, Deepwave利用了Pytorch可调用GPU进行反向传递和求导的特点, 具备更强大的运算性能.
实际上的话, Deepwave封装了大量的面向地球科学领域的可运用的算法, 它包括标量波动方程的2D正则和Born建模、弹性波动方程的正演、RTM、LSRTM、联合偏移反演等等.
本博客主要侧重于如何使用一些简单的内容, 侧重给偏计算机的研究人员调用.
如果想了解完整的代码和使用例子可通过其作者 (Alan Richardson) 的Github主页和文档说明来了解. 本文也将基于文档说明中的例子来延伸.

本篇的实验是基于Deepwave最新版本 (2024.5.14) v0.0.20
要获得这个版本你可以采用以下两种途径:

1. cmd 的 pip install deepwave 然后用pip list检查一下版本是否对应, 如果版本不对应, 可能你需要使用途径2
2. (针对windows) 到这里下载v0.0.20的zip版本, 然后解压到随便一个路径. 然后使用cmd的命令cd /d X:\XXX\XXX\deepwave-0.0.20 移动到这个路径, 执行pip install .来离线安装

1. Pytorch的自动求导机制

我们往往关注Pytorch在神经网络搭建上的便利而往往忽略了其最基本两项优秀的特性:

1. 调用GPU (借助Tensor)
2. 自动求导

为了搭建在GPU上, Pytorch创建一种不同于Numpy中的ndarray的全新数据结构: Tensor.
创建一个Tensor并将其搭建在GPU可以通过如下代码实现:

a = torch.arange(3)  								# [0, 1, 2] (搭载于CPU)
b = torch.ones(3, device=torch.device('cuda'))  	# [1, 1, 1] (搭载于GPU)
c = a.cuda()**2 + 2 * b  							# [2, 3, 6] (搭载于GPU)
d = c.cpu()  										# d = [2, 3, 6] (搭载于CPU), c仍然在GPU上

这种相互换载的手段需要注意的地方在于: 不同平台的变量之间是不能相互干涉计算的.

自动求导的关键是变量之间构成的关于梯度的链式关系.
例如下面这样的例子:

a = torch.arange(3.0, requires_grad=True)
b = (2 * a).sum()
print(b)											# tensor(6., grad_fn=<SumBackward0>)
b.backward()
print(a.grad_fn)									# None
print(b.grad)										# None
print(a.grad)										# tensor([2., 2., 2.])

在这里, 我们创建了一个Tensor变量a, 包含[0, 1, 2]并表明我们需要记录这个变量的梯度.
requires_grad=True是一个关键的声明, 声明之后对于这个变量的所有操作都将被追踪.
此外对于任意变量使用.requires_grad_()方法也可以手动开启追踪.
如果希望当前代码中如此标记的变量不被追踪, 可以用with torch.no_grad()包裹住目标代码来临时取消追踪.
如果想完全清除一个变量的所有"追踪"记录, 也可以直接使用.detach()

b是a通过求和操作得到的变量.
输出b的时候可以看见其grad_fn属性为<SumBackward0>.
这告知了两个信息: 首先, 因为a是被"追踪"的变量, 且b的运算包括了这个变量, 于是其也被"追踪";
其次, 因为b是对a求和得到的, 因此b包含对它的操作的伴随项a的引用"grad_fn=<SumBackward0>".
而a作为这个运算链的末尾叶子结点, 其grad_fn值为None. 这个运算链如下图

因为b触发了backward运算后, 以b为起点, 向它下方的所有节点传递, 计算所有requires_grad为True的变量的梯度grad.
但是b本身是没有梯度的, 因为这个节点并不是直接设置requires_grad为True的节点.
然后我们试图建立一个稍微复杂一点的求导链:

a = torch.arange(3.0, requires_grad=True)
b = (2 * a).sum()
c = torch.ones(3, requires_grad=True)
d = (a * c).sum() + b

print(d.grad_fn)					# <AddBackward0 object at 0x000001700D36AA20>
print(c.grad_fn)					# None				
print(b.grad_fn)					# <SumBackward0 object at 0x000001DB1E6DAA58>
print(a.grad_fn)					# None

d.backward()

print(d.grad)						# None
print(c.grad)						# tensor([0., 1., 2.])
print(b.grad)						# None
print(a.grad)						# tensor([3., 3., 3.])

从grad_fn可以发现, c与b作为叶子, 不是任何操作的结果, 因此没有对应的grad_fn值. 下图所示

而在d触发.backward()操作后, 这个链状结构将向下传递, 更新所有设置requires_grad为True的节点的grad值 (在不求导之前, 这些梯度值是None).
需要注意, 这里a变量的梯度是两条链路共同影响的结果.
如果我们不是对d使用.backward()而是对b使用, 那么最终c的grad就是None, 而a的梯度只有[2,2,2]

同时, 要注意一个问题, 假如我们在上个例子的基础上再重新声明d和b式子并对d使用.backward(), 那么我们会得到下面的结果:

...

b = (2 * a).sum()
d = (a * c).sum() + b
d.backward()

print(c.grad)				# tensor([0., 2., 4.])
print(a.grad)				# tensor([6., 6., 6.])

这个结果并不是出错了, 而是新的求导得到的梯度将叠加在旧梯度之上.
当你想计算许多数据样本的梯度, 但又没有足够的内存同时处理所有样本时, 这就很有用了, 因为这样可以分批处理并累积梯度.

Deepwave作者认为: 当要计算许多镜头的梯度时, 就会出现这种情况

如果要清空之前的梯度, 一种较为快捷的方式就是重新声明. 或者在使用了优化器的基础上, 对优化器使用zero_grad()
这个对于网络的batch计算是必要的, 因为当前batch计算得到的梯度并不会沿用到下次计算.

2. Pytorch实现最小化的机制

既然提到了优化器, 我们可以看下Pytorch的优化器可以做到什么程度.
PyTorch 提供了多种优化器, 你也可以使用梯度创建自己的优化器. 作为一个例子, 这里我们将介绍一个随机梯度下降法 (Stochastic Gradient Descent) 的最小化例子.

import torch

def f(x):
    return 3 * (torch.sin(5 * x) + 2 * torch.exp(x))

x_true = torch.tensor([0.123, 0.321])
y_true = f(x_true)

x = torch.zeros(2, requires_grad=True)
opt = torch.optim.SGD([x], lr=1e-4)
loss_fn = torch.nn.MSELoss()

for i in range(2000):
    opt.zero_grad()
    y = f(x)
    loss = loss_fn(y, y_true)
    print(y.grad_fn)						# <MulBackward0 object at 0x000001EE93DE2F28>
    print(loss.grad_fn)						# <MseLossBackward0 object at 0x000001EE93DE2F28>
    loss.backward()
    opt.step()

print(x)									# tensor([0.1230, 0.3210], requires_grad=True)

这里构建了一个曲线, 两个目标点和两个重叠于 $(2,f(2))$ 的初始解.
根本目标为: 从初始解出发来拟合目标点.

首先, 我们设置了一个"SGD优化器" (Pytorch的优化器中封装各种梯度下降的算法).
SGD具备两个参数, 一个是自变量参数$x$, 一个是下降的步长.

其次, 实现一轮梯度下降须知道一个函数关于这个自变量$x$的梯度. 因为

1. 理论上, 关于$x$的函数是得知$x$的下一步挪动方向/大小的基础.
2. 代码逻辑上, 计算梯度至少需要一个对$x$的运算链的结构, 仅有一个$x$肯定没法算梯度.

对于最小化问题, 这个函数往往用包含$f(x)$与目标$y_{\text{true}}$的MSE来表示.
其运算链如下:

只要我们在loss执行.backward()操作就可以顺着这条链路更新到唯一设置requires_grad为True的变量, 即$x$.

最后, 得到梯度后, 通过对优化器opt使用.step()方法来更新$x$.
因为优化器在声明时就与