Volve油田的3D地下声阻抗模型（通过PyLops的空间正则化反演估算得到）。致谢：本作品。
介绍“Python 很慢，”你听过这句话多少次？我想，很多人都听过很多次（可能是因为那些更熟悉其他编程语言或者对 Python 不是很熟悉的人）。就像任何流行的谚语一样，这句话里确实有些道理。那么，让我们试着从不同角度看为什么你总是听到这句话。
首先，如果我告诉你，“乌塞恩·博尔特速度慢”。你可能会觉得我疯了。乍一听确实如此，但如果我告诉你，“乌塞恩·博尔特今天速度慢，因为他这一周一直生病，而且他有一个裂开的趾甲”，你可能会更认同我的说法。同样的，对于Python来说，如果我告诉你，“我在用Python解决一个优化问题，但我不能使用NumPy，而且我必须使用列表来完成所有操作。Python运行起来会比较慢！”，我的论点或许就站得住脚了。但是现在是2024年，numpy是在2006年发明的，自那以后，Python科学社区已经开发了一套高性能的线性代数工具（例如，例如参见这篇文章，介绍性文章）。仅仅使用Python标准库来解决优化问题，而忽视这样一个生态系统，将是不理智的。
而且我们现在已经是2024年了，(Gen)AI革命已是我们世界的一部分有一段时间了。而且你知道吗，如今大多数GenAI工具都是用Python开发的，基于诸如Pytorch、JAX这样的库。甚至我们朋友Cerebras的世界上最快的AI芯片也可以直接通过Python使用（参见他们关于简单矩阵向量乘法的教程，以及用于训练GPT模型的示例代码）。那么，既然这些家伙都能直接用他们的超级昂贵的芯片从Python中操作，我们为什么不能用Python解决逆问题？
答案可能比你想象的还要明显：到目前为止，我（部分）误导了你。我在上一段提到的所有内容其实都不是真正的Python代码。Python实际上只是这些极其高效的C/C++/CUDA/Cerebras SDK库之上的一个薄层，这些库通过高级API暴露给Python程序员。
那么，如果你采用类似的策略，开发一个针对大规模逆问题的软件包会怎样呢？好消息是，你其实不需要这样做，因为 PyLops 现在已经为你做到了这一点！
开始PyLops 是一个开源的 Python 库，设计用于 线性算子相关的计算，旨在实现数值线性代数、信号处理、地球物理等科学领域的可扩展高效算法。— ChatGPT
首先，让我们考虑一个简单的示例逆问题，并向您展示什么是无矩阵方法以及如何利用PyLops来解决它。假设我们有一个信号 x(t) 和一个滤波器 h(t)，并且我们想从信号 y(t)= h(t) ∗ x(t) 中去除该滤波器，其中我使用 ∗ 表示卷积。这是一个在许多科学领域都非常流行的一种逆问题，称为去卷积（一种逆问题）。这个问题的一种简单的实现方法是简单地使用滤波器 (H) 构建一个托普利茨矩阵并进行求逆操作。这是因为这个问题是线性的，因此可以写成一个线性方程组：y = Hx。求逆操作后，它将应用于离散化的输出信号 (y)，以生成输入信号 (x) 的估计值。然而，随着问题规模增大，这种方法不再很好地适应。
在无矩阵的优化方法中，我们用一个更高效且更节省内存的方式来表示矩阵对向量（及其伴随向量）的作用。在这种情况下，我们可以轻松地使用 np.convolve 或 SciPy、CuPy、Jax 等中的类似方法。PyLops 的 Convolve1D 操作符正是这样做的，整个问题可以用几行代码轻松地在一个 CPU 上解决。
import numpy as np
import scipy as sp
from pylops.signalprocessing import Convolve1D
# 模型设置
nx = 11
x = np.zeros(nx)
x[nx//2] = 1
# 滤波器
nh = 3
h = np.ones(nh)
# 算子
H = Convolve1D(nx, h, offset=nh // 2)
# 数据 (@ 表示矩阵乘法，详情参见 PEP 465)
y = H @ x
# 逆向操作
xinv = sp.sparse.linalg.lsqr(H, y, iter_lim=100)[0]其中一个好处是PyLops与问题的数学表述非常接近。另一个优点是，只需对上面的代码做一些简单的修改（如#CHANGED所示），就能方便地在单个GPU上解决这个问题。值得注意的是，唯一重要的改动是将代码中的SciPy的LSQR替换为PyLops的LSQR，后者不支持GPU上的运算。
导入cupy作为cp
from pylops.signalprocessing import Convolve1D
from pylops.optimization.basic import lsqr
# 模型定义
nx = 11
x = cp.zeros(nx)
x [nx//2] = 1
# 滤波器定义
nh = 3
h = cp.ones(nh)
# 算子定义
H = Convolve1D(nx, h, offset=nh // 2)
# 数据定义
y = H @ x
# 逆运算
xinv = lsqr(H, y, niter=100)[0]到目前为止一切顺利，但这在PyLops的标准来看是一个很小的问题。接下来，我将考虑一个真正的科学问题，它建立在刚才介绍的反卷积概念之上。这个问题在地球物理处理中被称为地震反演，其目标是反演（或称为“反卷积”）地震图像，以获得地下岩石属性的估计（更具体地说是声阻抗模型(Z=cρ)，其中_c_表示速度，_ρ_表示密度）。
我们现在要解决的问题可以从物理上用以下关系来描述：y(t)= h(t) ∗ ∂ ₜ x(t)。请注意，这个问题与我们先前处理的问题非常相似，唯一的区别是输入信号上多了一个时间导数(∂ ₜ)。这并不是个问题，因为我们可以将其作为另一个滤波器来处理，将它和_h(t)_结合，或者对输入信号应用一系列操作（PyLops在处理这类问题上表现得非常好）。
然而，之前的玩具问题与当前问题之间存在一个重大差异：在实际应用中，人们通常感兴趣的是，一次性反演整个三维地震数据体（例如，720 × 401 等于 288,720 个一维向量——这就是我们接下来要解决的问题规模），这样可以引入正则化以促进平滑或块状结构等特性。作为一个例子，我们希望最小化如下的目标函数：
其中 Δ 通常称为拉普拉斯算子，在数学中指的是所有方向上的二阶偏导数之和，也就是。
好吧，大约288,000条痕迹（每条痕迹有800个样本）意味着我们的x和y向量大约有约2.31亿个样本。这相当于大约1GB的单精度浮点数据。这么大的矩阵求逆可能会非常耗时。使用无矩阵方法在单个CPU（即使是非常强大的）上解决这个问题也可能效率不高。
接下来，我们将再次依赖于PyLops及其四个层级的高性能抽象层，特别是其四个层级的高性能抽象。
- 仅使用单个CPU（主要使用
numpy数组和方法）- 仅使用单个GPU（主要使用
cupy数组和方法） - 多CPU（结合使用
numpy和mpi4py） - 多GPU（结合使用
cupy和mpi4py）
- 仅使用单个GPU（主要使用
所有用于重现以下结果的代码都可以在这个GitHub repository中找到，因此，我将在这里只关注一些关键代码行，这些代码行使用PyLops操作符和求解器对感兴趣的数据集进行地震逆时处理。最后，我会比较不同选项的性能，并讨论如何根据实际情况明智地选择最佳选项。但多GPU选项并不总是最佳选择——甚至绝对不是最经济的选择。
虽然这里的主要目标是让 PyLops 与 cupy 和 mpi4py 进行交互，交互，PyLops 同时它还支持与 PyTorch、Jax 和 PyTensor 的集成。
在开始之前，先简单介绍一下我将使用的输入数据集。这是一个尺寸为800×720×401（时间样本数 (nt) × 交叉线样本数 (nxl) × 内联样本数 (nil)）的地震数据集，该数据集作为Volve开放数据集的一部分公开提供。
我在工作站上完成了整个分析，那是一台挺厉害的机器。
- Intel(R) Xeon(R) Gold 6230R CPU @ 2.10GHz，拥有52个CPU核心（每个插槽26个核心，共计2个插槽）和264GB的内存。
- 2个GPU：1个Quadro RTX 4000（8GB）和1个英伟达 GeForce RTX 3090（24GB）。GPU：图形处理单元。
我们的基准解决方案将使用PyLops对地震反演的原始实现来获取，该实现运行在单个CPU上。在这种情况下，类似于示例，我们将完全在CPU上运行，如下方示意图所示。
单核模式下PyLops算子的前向传递示意图，credits: 自己的作品。
如果我们使用 PyLops 的 PoststackLinearModelling 操作符，上述目标函数可以很容易地按以下方式最小化：
# 操作符
PPop = PoststackLinearModelling(wav, nt0=nt, spatdims=(nxl, nil))
# 正则化器
Regop = Laplacian((nt, nxl, nil), axes=(1, 2), dtype=PPop.dtype)
# 逆解
niter_sr = 10 # 迭代次数 niter_sr
epsI_sr = 1e-4 # 阻尼系数 epsI_sr
epsR_sr = 0.1 # 空间正则化系数 epsR_sr
xinv = regularized_inversion(
PPop,
y.ravel(),
[Regop, ],
x0=x0.ravel(), # 初始猜测值 x0
epsRs=[epsR, ], # 正则化系数列表 epsRs=[epsR, ]
参数(iter_lim=niter_sr,
damp=epsI_sr,
show=True) # 参数字典，设置迭代限制、阻尼系数和显示信息
)[0]首先，我们使用尺寸为200×360×200的子体积，通过LSQR求解器进行10次迭代来解决这个问题。整个求逆过程总共大约需要45秒（每次迭代大约4.5秒）。
然后对整个数据范围进行反转：这个过程大约需要905秒（每次迭代需要91秒）。
单GPU地震逆时在PyLops中，从单CPU切换到单GPU非常简单（前提是你的问题能适应GPU的VRAM）；如下图所示，我们只需确保输入向量和操作符使用的任何数组都使用cupy而不是numpy。让我们看看如何修改上面的代码（更改部分用#CHANGED标记）。
单GPU模式下PyLops操作符前向传递的示意图。图源：本人作品。
# 运算符
PPop = PoststackLinearModelling(cp.asarray(wav), nt0=nt,
spatdims=(nxl, nil)) # 修改
# 正则化算子
Regop = Laplacian((nt, nxl, nil), axes=(1, 2), dtype=PPop.dtype)
# 逆
niter_sr = 10 # lsqr 迭代次数：
epsI_sr = 1e-4 # 阻尼系数
epsR_sr = 0.1 # 空间正则化：
xinv = regularized_inversion(
PPop,
cp.asarray(y).ravel(), # 修改
[Regop, ],
x0=cp.asarray(x0).ravel(), # 修改
epsRs=[epsR],
{'iter_lim': niter_sr,
'damp': epsI_sr,
'show': True}
)[0]我们再次在一个200 × 360 × 200大小的子体上进行反演处理，使用LSQR求解器进行了10次迭代。在这种情况下，整个反演处理仅需0.85秒（约0.09秒/次），在我的NVIDIA GeForce RTX 3090上运行，速度提升高达53倍，令人惊讶。
到目前为止还算顺利，但当我们尝试将问题扩展到整个地震数据体时，就遇到了大麻烦——OutOfMemoryError: 内存分配不足……如果你用的是GPU，肯定知道这是什么意思，游戏结束！
嗯，其实还是有的，总是有解决方法；目前我想出了三个解决方法。
- 手头还有点零钱，就出去买了一个更强大的GPU，显存更大。
- 我将问题拆分成更小的子问题，并独立解决它们，接受一些边缘效应的影响。
- 我切换到了多GPU环境。
在这三个选择中，只有最后一个是最理想的；但这意味着，我需要重新思考我的代码设计，因为我不能像之前那样直接从numpy切换到cupy。不用担心，PyLops又一次来救场了。在我们解决多GPU环境下的反演问题之前，退一步，也考虑一下多CPU环境下的情况。
多核地震反演可能再次有远离单CPU设置的两个原因：
- 问题太大，根本放不进CPU内存；
- 问题可能勉强能塞进CPU内存里，但反转过程极其缓慢，让人难以忍受。
在这个例子中我们遇到的情况是，在反转整个地震数据体时所面对的情形。
从现在开始，我们要介绍一个新的东西：PyLops-MPI，
PyLops-MPI 支持 MPI（消息传递接口），从而支持大规模线性算子问题的分布式内存并行处理。它专门设计用于处理单机无法处理的大规模线性算子问题，利用多节点集群或超级计算机的计算资源。
我假设读者熟悉MPI（以及MPI4py），因此我会跳过所有用于MPI驱动Python脚本的样板代码。我将再次聚焦于关键的PyLops/PyLops-MPI调用，这些调用使我们能够在多个CPU上完成地震反演任务。
对于这类问题，行动计划很简单。由于建模算子在追踪级别上是天然并行的，我们将把追踪平均分配到所有可用的 rank 上，并让每个 rank 对其一组追踪应用建模算子。为了简化分配，我将沿着内联方向上分割追踪（这样每个 rank 将处理一组内联的所有追踪）。在 PyLops 术语中，整体算子是一个分块对角算子，因为每个独立的算子作用于一组追踪，并且所有算子彼此完全独立。PyLops-MPI 提供了一个 MPIBlockDiag 算子，它接受一个或多个算子在每个 rank 上，并且在所有 rank 之间表现得就像一个全局的分块对角算子。再次，我们可以在下面的图中再次看到这一点。
PyLops操作符在多CPU模式下的前向传递示意图，虚线表示总向量（x/y）与其在各个MPI进程中的部分之间的逻辑连接。注：本图由本人绘制。
但是，对于拉普拉斯算子，需要一个更细化的策略，因为它要求在内联（即分布式）轴上应用二阶导数，我不会深入其实现细节。感兴趣的读者可以查阅我们的文档页面。最后，求解器本身也需要执行多个跨进程的操作（例如点积或范数）。为此，PyLops-MPI 提供了其 [分布式数组](https://pylops.github.io/pylops-mpi/api/generated/pylops_mpi.DistributedArray.html#pylops_mpi.DistributedArray)（可以将其视为在多个进程中分布的增强版 NumPy 数组）。现在我们来看看代码是如何变化的：
# 选择如何将ilines划分到各个rank
nil_rank = local_split((nil, ), comm, Partition.SCATTER, 0)
nil_ranks = np.concatenate(comm.allgather(nil_rank))
ilin_rank = np.insert(np.cumsum(nil_ranks)[:-1] , 0, 0)[rank]
ilend_rank = np.cumsum(nil_ranks)[rank]
ilines_rank = ilines[ilin_rank:ilend_rank]
# 提取感兴趣的数据和初始模型部分
# （在实际操作中，这部分数据通常是直接从文件中读取的）
y = y[ilin_rank:ilend_rank]
x0 = x[ilin_rank:ilend_rank]
# 分布式数据集
y_dist = pylops_mpi.DistributedArray(
global_shape=nil * nxl * nt,
local_shapes=[(nil_r * nxl * nt,) for nil_r in nil_ranks],
dtype=np.float32)
y_dist[:] = y.flatten().astype(np.float32)
# 分布式初始估计
x0_dist = pylops_mpi.DistributedArray(
global_shape=nil * nxl * nt,
local_shapes=[(nil_r * nxl * nt,) for nil_r in nil_ranks],
dtype=np.float32)
x0_dist[:] = x0.flatten().astype(np.float32)
# 操作符（Operator）
PPop = PoststackLinearModelling(wav, nt0=nt,
spatdims=(nil_rank[0], nxl))
Top = Transpose((nil_rank[0], nxl, nt), (2, 0, 1))
BDiag = pylops_mpi.basicoperators.MPIBlockDiag(ops=[Top.H @ PPop @ Top, ])
# 正则化算子
Regop = pylops_mpi.MPILaplacian(dims=(nil, nxl, nt),
axes=(0, 1, 2),
dtype=BDiag.dtype)
# 逆操作
niter_sr = 10 # lsqr迭代次数
epsR_sr = 0.1 # 空间正则化
StackOp = pylops_mpi.MPIStackedVStack([BDiag, np.sqrt(epsR_sr) * LapOp])
y0_dist = pylops_mpi.DistributedArray(global_shape=nil * nxl * nt)
y0_dist[:] = 0.
ystack_dist = pylops_mpi.StackedDistributedArray([y_dist, y0_dist])
xinv = pylops_mpi.optimization.basic.cgls(StackOp, ystack_dist,
x0=x0_dist,
niter=niter_sr,
show=True)[0]注意，在这种情况下，我们使用了额外的 Transpose 操作符来重新组织输入数组，然后再将其输入到 PoststackLinearModelling 操作符中。这是必要的，因为 [DistributedArray](https://pylops.github.io/pylops-mpi/api/generated/pylops_mpi.DistributedArray.html#pylops_mpi.DistributedArray) 只能在最慢的轴（这里为测线（inline））上进行分割，而 PoststackLinearModelling 操作符需要输入的数组的第一维度是时间。除了这一小的实现细节外，代码看起来相当直观。每个进程首先确定要处理的测线（inline）组，并提取相关部分的数据和初始模型。然后将数据和模型都嵌入到 DistributedArray 中。将 Transpose-PoststackLinearModelling-Transpose 这个链路包装到一个 MPIBlockDiag 操作符中。
我们的逆向代码与我们迄今为止看到的有所不同；事实上，为了求解上述方程中的正则化逆问题，会创建一个扩充的方程组——请注意，在PyLops中，这是在regularized_inversion方法内部自动完成的。让我们来看看在纸上看来是怎么样的，以及实际上我们是如何实现这一点的。如果你熟悉逆问题，你会立刻发现上面的目标函数可以等价地写成：
但既然这两种表达是相同的，为什么这会有关系呢？答案是抽象。如果我们能创建一个抽象的、通用的表达，我们就可以始终使用同一个解算器。首先，我们通过将数据y与一个零向量堆叠来生成增强型数据。
# 增强数据
y0_dist = pylops_mpi.DistributedArray(global_shape=nil * nxl * nt) # nil 可能代表空值或零值
y0_dist[:] = 0. # 将 y0_dist 的所有元素初始化为 0
ystack_dist = pylops_mpi.StackedDistributedArray([y_dist, y0_dist])然后也对操作员这样做
# 扩展操作符
# PoststackLinearModelling 是用于后叠加线性建模的函数
PPop = PoststackLinearModelling(wav, nt0=nt,
spatdims=(nil_rank[0], nxl))
# Transpose 用于矩阵转置操作
Top = Transpose((nil_rank[0], nxl, nt), (2, 0, 1))
BDiag = pylops_mpi.basicoperators.MPIBlockDiag(ops=[Top.H @ PPop @ Top, ])
# MPILaplacian 是用于多进程Laplacian操作的函数
Regop = pylops_mpi.MPILaplacian(dims=(nil, nxl, nt),
axes=(0, 1, 2),
dtype=BDiag.dtype)
# MPIStackedVStack 是用于多进程堆叠操作的函数
StackOp = pylops_mpi.MPIStackedVStack([BDiag, np.sqrt(epsR_sr) * LapOp])
# "@" 表示矩阵乘法操作最后，调用了 CGLS 解算器的 PyLops-MPI 实现——为什么不使用 LSQR 呢？很简单，PyLops-MPI 还不支持这个解算器，而 CGLS 和 LSQR 实际上很接近。
这时，我们已经准备好将问题分到各个内联中了——但该怎么分呢？首先，让我们先给出一个简单的答案：
导出环境变量OMP_NUM_THREADS, MKL_NUM_THREADS和NUMBA_NUM_THREADS为4，然后使用mpiexec并行执行Poststack_Volve.py脚本8次。export OMP_NUM_THREADS=4;
export MKL_NUM_THREADS=4;
export NUMBA_NUM_THREADS=4;
mpiexec -n 8 python Poststack_Volve.py
在那，我确保创建各种 *_NUM_THREADS 环境变量（例如 OMP_NUM_THREADS），这样就能确保 numpy 及其他利用多线程的 Python 库不会尝试使用我全部的 52 个线程。
最难的部分是：我该选择什么值作为 -n X。例如，我可以选择 n=2 将我的逆操作分布在2个等级上，或者选择 n=4 将其分布在4个等级上。原则上来说，越多越好，但这也引入了更多的通信成本。我测试了以下4种情况来解决整个问题：
*_NUM_THREADS=16，n=2：878秒（~每迭代88秒）*_NUM_THREADS=8，n=4：492秒（~每迭代49秒）*_NUM_THREADS=4，n=8：310秒（~每迭代31秒）*_NUM_THREADS=2，n=16：264秒（~每迭代26秒）
所以在所有情况下，倒置所需时间都比单核CPU的情况要短；在最佳情况下，我获得了约×3的速度提升（公平地说，我在多核CPU情况下将*_NUM_THREADS * n设为了32，而在单核CPU情况下我设*_NUM_THREADS为32）。需要注意的是，对于这个问题来说，最好增加rank的数量并减少线程数——但这可能不适用于所有问题，例如当操作本身很好地利用了线程并行性时情况就不同了。
最后，我们来谈谈多GPU的情况。我先来看看理想的情况是怎样的，我有足够的GPU，它们的总显存可以支持我想要解决的问题规模。
PyLops操作符在多GPU模式下前向传递的示意图。虚线表示总向量（x/y）与其在多个MPI进程间分布的各个部分之间的逻辑连接。注：本图由本人绘制。
在这种情形下，修改上述代码非常简单，因为我们可以直接使用MPI4py中的CUDA-aware MPI支持（如#CHANGED所示更改的部分，……则表示未修改的代码部分）。
# 选择如何拆分等级
...
# 提取部分数据和感兴趣的相关初始模型
# （在实践中这将很可能会直接从文件中读取）
...
# 分布式数据
y_dist = pylops_mpi.DistributedArray(
global_shape=nil * nxl * nt,
local_shapes=[(nil_r * nxl * nt,) for nil_r in nil_ranks],
dtype=np.float32, engine='cupy') # 注释
y_dist[:] = cp.asarray(y.flatten()).astype(np.float32) # 注释
# 分布式初始估计
x0_dist = pylops_mpi.DistributedArray(
global_shape=nil * nxl * nt,
local_shapes=[(nil_r * nxl * nt,) for nil_r in nil_ranks],
dtype=np.float32, engine='cupy') # 注释
x0_dist[:] = cp.asarray(x0).flatten().astype(np.float32) # 注释
# 算子
PPop = PoststackLinearModelling(cp.asarray(wav), nt0=nt, # 注释
spatdims=(nil_rank[0], nxl))
...
# 正则化
...
# 求解逆问题
...然而，由于我的硬件限制条件（其中一个GPU只有8GB的VRAM），我只能对子体进行地震反演；在这种情形下，整体反演过程需要3.65秒（每迭代一次大约0.4秒）。
我们之前说过，如果无法在单个GPU上解决整个问题，我们可以利用多GPU设置。但似乎即使在这种情况下，我们有时也会因为可用资源不足而内存溢出。让我们看看最后一个选项，即使用混合CPU-GPU设置。换句话说，我们的问题被分摊到多个rank上，而数据（和模型）存储在CPU的RAM中。然而，每当我们要计算建模算子时，会将位于单个rank中的部分迹线移动到GPU上，执行Transpose-PoststackLinearModelling-Transpose操作链，然后将输出移回CPU。这样，我们不需要在整个过程中将所有数据（和模型）都保存在GPU的总VRAM中，而只需要在需要时将小部分数据移动到GPU进行操作。
这种方法看起来不错且可扩展，但在你感到失望之前，我想先说一句：通过这样做，在数据和模型能够自信地放置在GPU(s)的VRAM中时，我们永远无法超越之前提到的单GPU或多GPU方法的性能；然而，我们可以解决更大规模的问题，并且（希望）在多CPU设置中仍能获得一些加速。那么我们如何实践这个想法呢：我们已经有了几乎所有需要的部件，只缺少一个简单的操作符。这个操作符的前向传递会将位于CPU上的输入（一个numpy数组）移动到GPU（并返回一个cupy数组），而它的伴随操作则正好相反。这个操作符叫做[ToCupy](https://pylops.readthedocs.io/en/latest/api/generated/pylops.ToCupy.html#pylops.ToCupy)，是PyLops库最近新增的功能之一。
简单来说（这里为了简化说明，只考虑一个CPU+GPU），同样的操作会在每个CPU+GPU上重复进行，我们所做的如下。
PyLops操作在多GPU模式下的前向传播示意图，其中输入向量分块依次处理。图片来源：自行绘制。
最后让我们看看如何在多GPU示例中应用这一点（再次强调……表示与多CPU代码相比没有任何改动）
# 选择如何将 ilines 分割成不同的等级
...
# 提取部分数据和感兴趣的初始模型
# （在实际情况中，这部分内容可能直接从文件中读取）
...
# 分布式数据
...
# 分布式初始猜测
...
# 算子
# 将 nil_rank[0] 分成若干组，这些组将被 GPU 依次处理
ngroups = 2
nil_pass = nil_rank[0] // ngroups
nil_pass_in = np.arange(0, nil_rank[0], nil_pass)
nil_pass_end = nil_pass_in + nil_pass
nil_pass_end[-1] = nil_rank[0]
BDiags = []
for nil_i, nil_e in zip(nil_pass_in, nil_pass_end):
PPop = PoststackLinearModelling(cp.asarray(wav), nt0=nt,
spatdims=(nil_e - nil_i, nxl))
Top = Transpose((nil_e - nil_i, nxl, nt), (2, 0, 1), dtype=np.float32)
TCop = ToCupy(Top.dims, dtype=np.float32)
TCop1 = ToCupy(Top.dimsd, dtype=np.float32)
BDiag = TCop1.H @ Top.H @ PPop @ Top @ TCop
BDiags.append(BDiag)
BDiags = BlockDiag(BDiags, forceflat=True)
BDiag = pylops_mpi.basicoperators.MPIBlockDiag(ops=[BDiags, ])
# 正则化项
...
# 逆操作
...关键在于创建一个_2级_的块对角操作，其中第一级是在GPU层面（模型的那部分位于CPU的RAM上，被切分后依次发送到GPU处理），第二级是在CPU层面（整个模型被切分并在各个进程中分布）。数据处理也遵循同样的原则，如上图所示。
你可能注意到了这一行有些晦涩的表达：nil_pass = nil_rank[0] // ngroups。因为不能将分配给某个特定排名（nil_rank[0]）的整个迹线块直接传递给GPU，所以我们需要决定将它分成多少部分（确保nil_rank[0] // ngroups个迹线可以放入我们的GPU显存中）。类似于我们在多CPU示例中选择的CPU数量n，选择ngroups可能需要一些试错来找到最适合你硬件的最优配置。
测试了以下两种情况以针对整个问题：
ngroups=2：441秒（约44秒/次）ngroups=4：441秒（约44秒/次）
在两种情况下，反演所需的时间都少于4个rank的多CPU情况，但多于使用4个rank以上的情况。这一结果可能是因为每次应用算子时，主机到设备和设备到主机的数据传输引入了额外的开销。此外，由于Laplacian算子并不是完全可分的（即不能写成块对角算子，因为沿着分块的维度的导数，此处为内联的组），我们的正则化项仍然完全在CPU上运行，而只有PoststackLinearModelling在GPU上运行。
要看到真正的性能提升，我们可能需要考虑一个不同的问题，其中操作的成本远远超过了数据传输操作的成本（这种情况对于PoststackLinearModelling操作来说可能性不大）。或者，我们可以稍微放宽正则化项，并将其拆分到各个内联组中，从而得到以下问题（对于ngroups=2）：
这里H1和H2是PoststackLinearModelling操作符，仅处理一组内联（inline），y1和y2包含这些内联的数据部分。从实现的角度来说，主要的变化仅仅是创建这个操作符。
# 操作符和数据
# 将 nil_rank[0] 分成若干组（依次由 GPU 处理）
ngroups = 2
nil_pass = nil_rank[0] // ngroups
nil_pass_in = np.arange(0, nil_rank[0], nil_pass)
nil_pass_end = nil_pass_in + nil_pass
nil_pass_end[-1] = nil_rank[0]
BDiags = []
ddiags = []
for nil_i, nil_e in zip(nil_pass_in, nil_pass_end):
PPop = PoststackLinearModelling(cp.asarray(wav), nt0=nt,
spatdims=(nil_e - nil_i, nxl))
Top = Transpose((nil_e - nil_i, nxl, nt), (2, 0, 1), dtype=np.float32)
TCop = ToCupy(Top.dims, dtype=np.float32)
TCop1 = ToCupy(2 * Top.shape[0], dtype=np.float32)
LapOp = Laplacian(dims=(nil_e - nil_i, nxl, nt),
axes=(0, 1, 2),
dtype=np.float32)
BDiag = TCop1.H @ VStack([Top.H @ PPop @ Top,
np.sqrt(epsR_sr) * LapOp]) @ TCop
BDiags.append(BDiag)
ddiags.append(np.hstack([
d[nil_i:nil_e].flatten().astype(np.float32),
np.zeros_like(m0[nil_i:nil_e].flatten()).astype(np.float32)]))
BDiags = BlockDiag(BDiags, forceflat=True)
BDiag = pylops_mpi.basicoperators.MPIBlockDiag(ops=[BDiags, ])
# 使用正则化方程进行带正则化的逆解
dstack_dist = pylops_mpi.DistributedArray(
global_shape=2 * nil * nxl * nt,
local_shapes=[(2 * nil_r * nxl * nt,) for nil_r in nil_ranks],
dtype=np.float32)
dstack_dist[:] = np.concatenate(ddiags)
# 逆解
...使用这种宽松的公式来运行逆向过程可以得到一个非常接近理想的情形的解，同时性能也有了明显的提升，因为现在Laplacian运算也在GPU上完全运行了，从441秒（约44秒/迭代）减少至184秒至194秒（约18秒至19秒/迭代），相较于最佳多CPU结果还要快。
总之，我提供了一个表格，总结了针对不同实验使用的小数据集和大数据集中的逆向变换耗时。请注意，在多CPU/GPU设置中，我没有对小数据集进行逆向变换，因为这会导致性能不佳（-表示未使用给定参数集进行逆向变换）。此外，内存不足错误指的是问题大小和GPU显存的组合导致我们无法解决问题的情况。
继续前行。如果你仍然在这里（做得不错！），你可能在想下一步是什么——无论是对你（用户）来说还是对我们（开发小组）来说。
让我们从你开始谈起：希望这个教程能激励你尝试在你领域中通常遇到的一些逆问题上使用PyLops。你可能会发现你的问题适用于单个*场景，希望你也能像我们在示例中从CPU迁移到GPU时一样获得加速效果。或者你可能会发现你的问题更像这个教程中的那种问题，其中多*场景才是解决现实中的大规模问题的唯一途径。
谈到开发时，我们已经发现（尤其是在本教程接近尾声时），仍然有改进的空间。特别是在多GPU场景中，当可用GPU的累计VRAM不足以容纳当前的问题时。一些改进的想法有：
- 通过 CUDA 流重叠数据移动和计算；
- 实现通过 GPU 加速的导数运算符，这些运算符顺序地在模型和数据向量的边界区域（halo）上运行（就像我们在 PyLops-MPI 中使用
[MPIFirstDerivative](https://pylops.github.io/pylops-mpi/api/generated/pylops_mpi.basicoperators.MPIFirstDerivative.html)和[MPISecondDerivative](https://pylops.github.io/pylops-mpi/api/generated/pylops_mpi.basicoperators.MPISecondDerivative.html)这样的操作符一样）。
我要感谢所有为PyLops的发展做出贡献的人，还要感谢Google暑期代码项目支持PyLops-MPI的发展。我也很感谢Equinor及其合作伙伴共同发布了公开的Volve数据集，本教程使用了该数据集。最后，还要感谢Carlos da Costa审阅了本文的早期版本。
资源列表- 本教程中使用的笔记本和脚本，
- PyLops参考手册
- PyLops-MPI参考手册
共同学习，写下你的评论
评论加载中...
作者其他优质文章
