2023 7.24~7.30 周报 (VelocityGAN)

0 上周回顾

上周读完了论文《Deep-Learning Full-Waveform Inversion Using Seismic Migration Images》.
收集了一些idea.

1 本周计划

这周继续深入阅读一篇DL-FWI论文:
《Data-Driven Seismic Waveform Inversion: A Study on the Robustness and Generalization》
并且通过已有代码来分析实现细节.
最终复现.

1.1 论文背景

这篇论文是《Data-Driven Seismic Waveform Inversion: A Study on the Robustness and Generalization》. 它发布于2020年, 所在期刊是《IEEE TRANSACTIONS ON GEOSCIENCE AND REMOTE SENSING》, 即TGRS.
这篇论文与经典论文InversionNet的有着共同的第二作者, 我有理由怀疑它们应该是共同的一个研究团队. 而且这篇论文的实验风格, 数据展示方式都与InversionNet非常相似. 而且通过阅读逐渐发现: 这篇论文有沿用InversionNet的许多结构, 不如说: InversionNet + Discriminator 构成的DGAN结构.
这篇论文中, 作者将它的结构称之为VelocityGAN.

2 完成情况

2.1 背景简述

按照惯例, 作者介绍了传统FWI的弊端, 但是作者的侧重点可谓别出心裁, 因为它着重强调了传统FWI在正则优化上的改进.

传统FWI方法相比于更简单的走时反演有更加精确的分辨率和精度, 但是FWI也具有更高的难度, 因为它的正向算子是非线性的, 模拟量大且问题具有不适定性, 没有唯一的解决方案. (三方面: 不适定性, 周期跳跃, 高计算成本)
于是近几年人们提出正则化的优化方法:

• Tikhonov正则化
• 全变分正则化
• 高阶正则化
• 基于先验的正则化

来解决FWI存在的多数问题, 但是这些方法都还是在物理方法的框架下架构的. 他们会被物理先验的精确性影响, 而且开销巨大.
虽然局限存在, 但正则化的优秀之处仍然不可无视, 因此在进入数据驱动FWI的研究时代下, 仍然有些很好的正则约束在沿用, 只不过多了很多"深度学习"的味道.
而神经网络结构存在一种非常特殊的正则约束, 那就是对抗约束, 即通过构造判别器 (Discriminator) 来与固有的网络 (Generator) 生成的数据对抗, 从而约束固有网络. 这就是GAN.
虽然GAN在之前有试着FWI领域进行融合 (论文: Generative Adversarial Networks for Model Order Reduction in Seismic Full-Waveform Inversion), 但是并不及VelocityGAN的名气, 因为相比之下, VelocityGAN有着较好的实验和论证, 包括损失函数的设计. 以及, VelocityGAN的数据集更广泛, 覆盖了合成和真实数据 (SEG), 同时也运用了迁移学习 (迁移到断层).

2.2 网络结构: 生成器 (Generator)

图的左侧是VelocityGAN中所提到的生成器网络的参数, 而右侧是InversionNet提供的网络参数.
实际上仔细观察可以发现, 两个使用的是同一套架构, 甚至在数据上也几乎是同一套.
InversionNet论文中直接演示的架构的信息里, 地震图像是$6\times 1000\times 32$, 速度模型用了两套, 分别是平坦结构的$100\times 100$和曲面结构的$100\times 150$ (2019年版)
VelocityGAN直接演示的架构的信息里, 地震图像也是$6\times 1000\times 32$, 速度模型只用了曲面数据, 大小也是$100\times 150$ (2020年版), 但是额外添加了SEG的数据进行真实数据测试.
InversionNet在解码器部分设置了$h\times w$来表示其对于速度模型的灵活控制, 只要设置$h=10$和$w=7$那么就可以最终获得$160\times 112$的图像. 再通过截取即可获得$100\times 150$的速度模型.

2.3 网络结构: 判别器 (Discriminator)

作者设置了一个较为浅层的CNN网络来表示判别器网络.

In particular, it consists of five convolution blocks, a global average pooling layer, and fully connected layers. Each convolutional block involves a combination of Convolutional, BatchNormalization, LeakyReLU, and MaxPooling layers.

在论文中, 他描述的判别器网络由五个卷积块, 一个全局平均池化层和全连接层组成, 但是并没有给出具体的结构参数. 于是我根据猜测绘制了其如下的图像.
作者描述的每个卷积块都包含了最大池化的操作, 其目的肯定是为了缩小图像尺寸, 由此可以推断其经历了至少5次尺寸缩小.
而通道数是我的一种猜测, 因为其的生成器也采用了类似的通道数规律设计.
最后的全局平均池化就是将最后512通道的每张4*5的图像都在图像内部求平均值, 即20个数值的平均, 由此会得到512个平均值. 这些平均值作为全连接神经网络的输入, 最后通过全连接网络, 压缩得到4个网络节点.

We apply the “PatchGAN” classifier [18] in the discriminator to capture local style statistics. We set the patch size as 4 and calculate the mean loss value of all patches in an image. In our data set, velocity maps can be a rather different one from the others (such as different tilting angles, layer thicknesses, and layer velocities). However, if we focus only on the local information, their geological faults and interfaces share some similarities (the geological faults and interfaces always have a drastic change in velocity). Therefore, “PatchGAN” is more suitable than “GlobalGAN” for our task.

这4个节点都是判别器的输出, 依据作者在文章中提到的PatchGAN思路, 这四个输出值反映了速度模型图像中 “四个区域” 的真实性判别,每个区域给出0.0~1.0的评分, 越高, 证明判别器认为它是real的概率越高, fake的概率越低.
这只是四个区域的判别, 最终, 四个区域的值求平均, 即是速度模型本身的真实性判别. (详情可百度"PatchGAN")

2.4 损失函数: Loss

VelocityGAN的损失函数使用了WGAN-GP的损失函数, 这部分百度倒都是有解释, 这里我为了方便总结, 我也捋一捋.
首先要明确, 对于GAN类的损失函数, 一般来说我们是先看 判别器 的损失函数, 因为程序里面, 我们确实要先告诉判别器什么是fake, 什么是real. 而往往来说, 喂入判别器的real就是生成器的真实标签 (Ground truth), 而fake是生成器预测的标签. 对于FWI, 标签即速度模型.

判别器Loss (如果是纯WGAN):
$$L_d=\underset{\tilde{x}\sim {\mathbb{P}}_g}{\mathbb{E}}D(\tilde{x})-\underset{x\sim {\mathbb{P}}_r}{\mathbb{E}}D(x)$$
这里是WGAN的判别器网络的loss, 这里的$\tilde{x}$表示预测的速度模型, 也就是fake图像; $x$是真实速度模型, 即real图像. 这里$\mathbb{E}$我通过原文分析可能是一种求平均值的算子, 即$\mathbb{E}(\text{images})$会将所囊括的图像的每个像素求平均.
此外, 这里的${\mathbb{P}}_g$表示生成器生成的速度模型的分布, 而${\mathbb{P}}_r$表示真实速度模型的分布. 言下之意, 投入到函数$\mathbb{E}(\text{images})$中的将会是一批符合特定分布的图像, 我们将会求它们每个的平均, 而后将会再度将每个图的平均值集结起来平均一次.
WGAN没用采用$\log (\cdot )$的方案, 同时因为Wasserstein距离的特点, 即它是计算两种概率分布之间的差, 因此它是可以为负的. 没错, 这个loss可以为负.
传统的WGAN还采用了Weight clipping, 也就是限制了网络权值的上下限$[-c,+c]$, 从而简介约束了梯度变化的不要太猛, 这样我们的$D(\tilde{x})$本身的值也会约束在一定范围了, 因为$D(\tilde{x})$展开后也是包含$w$的函数, 这样即使不通过log等手段, WGAN的loss也不会有梯度的大规模变化.
而WGAN-GP在右侧添加了一个约束的部分:
$$L_d=\underset{\tilde{x}\sim {\mathbb{P}}_g}{\mathbb{E}}D(\tilde{x})-\underset{x\sim {\mathbb{P}}_r}{\mathbb{E}}D(x)+\lambda \underset{\hat{x}\sim {\mathbb{P}}_{\hat{x}}}{\mathbb{E}}\left[{\left({\Vert {\nabla }_{\hat{x}}D(\hat{x})\Vert }_2-1\right)}^2\right]$$
它在添加了这部分至于取消了WGAN原本的Weight clipping, 原因很简单:

权值裁剪限制了网络的表现能力。由于网络权值被限制在了固定的范围内，神经网络很难再模拟出那些复杂的函数，而只能产生一些比较简单的函数

因此WGAN-GP提出了采用梯度惩罚项 (GP: Gradient penalty), 即公式右侧冒出的那部分, 来代替Weight clipping.
梯度惩罚是基于fake和real之间的混合图像$\hat{x}$在判别器之后的梯度二范数来约束的.
这个混合图像是基于一个随机参数来调整real和fake图像的占比: $\hat{x}=ϵ\tilde{x}+(1-ϵ)x,ϵ\sim U[0,1]$, 这是一种差值过程, 保证在real和fake图像之间的Lipschitz连续
这种梯度二范数在接近1的时候影响最小.

然后我们看下关于生成器网络的loss设计:
$$L_g=-\underset{\tilde{x}\sim {\mathbb{P}}_g}{\mathbb{E}}D(\tilde{x})+\frac{{\lambda }_1}{w\cdot h}\sum _{i=1}^w\sum _{j=1}^h\mid \tilde{v}(i,j)-v(i,j)\mid +\frac{{\lambda }_2}{w\cdot h}\sum _{i=1}^w\sum _{j=1}^h{\left(\tilde{v}(i,j)-v(i,j)\right)}^2$$
这个loss的后面两个组件很简单, 其实就是$l_1$和$l_2$的常规损失函数, 生成器网络本身就是一个很普通的端到端FWI.
而第一个部分的loss组件是GAN类loss特有的对抗性约束, 这是强化一般生成网络的很重要的约束, 也是VelocityGAN的性能效果优于InversionNet的主要改进.
因为通过之前对于生成器网络的损失梯度传递 (BackPropagation), 判别器网络目前已经优化了对于真实标签和预测标签的判定, 这时, 如果生成器能提供一个足够以假乱真的标签: 一来说明我预测得完美, 二来也能降低loss中第一个对抗性约束项的值.
而且这两点会相互激励, 因为对抗性约束一旦高了, loss的特性可得, 它会通过反向传递来进一步降低所有loss组件的值来激励对抗性约束达到更小的值. 这种过程会进一步促进生成器拟合得更好, 将预测标签进一步向真实标签靠近.

2.5 OpenFWI中的VelocityGAN与它的核心代码

OpenFWI这篇论文 (OpenFWI: Large-Scale Multi-Structural Benchmark Datasets for Seismic Full Waveform Inversion) 是2022年多个作者联合发布的关于大量合成数据集开源化的综述性论文. 一方面, 这篇论文提供了超过的开源数据; 另一方面, 他开源了某些合作作者参与过的DL-FWI有关的论文, 这里面包括了InversionNet, VelocityGAN, UPFWI和InversionNet3D.
而这些源码中的网络结构和OpenFWI论文中的训练参数, 相比于原论文都是有着细微不同.
因为OpenFWI发布于2022年, 因此这种改进后的版本的算法, 我愿称之为2022版本.

而2022版本的VelocityGAN相比于2020版(原版)有如下变化:

• 生成器网络采用了2022版的InversionNet.
• 输出的速度模型大小变为70*70.
• 输入的地震记录大小变为1000*70, 共5炮构成.
• 判别器网络采用新结构, 最后的PatchGAN似乎被取消了, 同时全连接组件被代替
• 训练的BatchSize设定50→64, 取消的学习率的梯度线性下降到0, 公布的训练轮次为480, 其中1次epoch更新1次判别器, 3次epoch更新1次生成器.

其中最大变化是判别器网络的结构, 基于已知代码, 我们可以绘制出一个确定的网络结构图:

2022版的判别器网络在结构更趋向于FCN的风格, 删掉了全连接, 模仿编码器的思路, 将速度模型压缩为一个长度为1*1的图像, 且通道数为1, 说人话就是直接压缩为一个数字.
自然, 这个版本抛弃了PatchGAN的原理. 最后只是对于网络同个batch内的所有输出的值取平均值操作即可.

下面给出这部分的代码设计: (Python)

2.5.1 判别器loss

class Wasserstein_GP(nn.Module):
    def __init__(self, device, lambda_gp):
        super(Wasserstein_GP, self).__init__()
        self.device = device
        self.lambda_gp = lambda_gp

    def forward(self, real, fake, model):
        gradient_penalty = self.compute_gradient_penalty(model, real, fake)
        loss_real = torch.mean(model(real))
        loss_fake = torch.mean(model(fake))
        loss = -loss_real + loss_fake + gradient_penalty * self.lambda_gp
        return loss, loss_real-loss_fake, gradient_penalty

    def compute_gradient_penalty(self, model, real_samples, fake_samples):
    	# real_samples.size(0) 是batch数, 我们要为不同的batch设置不同的随机数
        alpha = torch.rand(real_samples.size(0), 1, 1, 1, device=self.device)
        interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)
        d_interpolates = model(interpolates)
        gradients = autograd.grad(
            outputs=d_interpolates,
            inputs=interpolates,
            grad_outputs=torch.ones(real_samples.size(0), d_interpolates.size(1)).to(self.device),
            create_graph=True,
            retain_graph=True,
            only_inputs=True,
        )[0]
        # 得到的gradients与输入维度一致, 即(64x1x70x70), 通过本行代码之后, 后续三维被压缩为4900, 即(64x4900)
        gradients = gradients.view(gradients.size(0), -1)
        # dim = 1 是对于矩阵的每行求2范数, 也就是将每行视为一个向量, 求这个向量的2范数, 把向量中每个元素平方后相加再开根号.
        gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()                                                 
        # 所以(64x4900)的矩阵, 每行4900长度的向量都变为单个值, 即最后矩阵变为(64x1)                                                                                                     
        return gradient_penalty

代码中的model函数对应判别器网络的对象, 可以理解为$D(\cdot )$; real就是速度模型标签, 可以理解为$x$; fake就是生成器网络预测的速度模型, 可以理解为$\tilde{x}$, torch.mean()执行的过程与$\mathbb{E}(\cdot )$是一致的.

alpha对应公式中的随机算子$ϵ$.
这里范围的real_samples和fake_samples都是成批次输入到网络中的图像数据, 且与地震数据对应, 单个数据本身就会有三个通道, 加上批次维度共有四个通道. 就OpenFWI版本来看, 维度将会是(batch_size $\times 1\times 70\times 70$)
需要注意, $ϵ$对于不同批次的数据是随机的, 而并非一个批次内的速度模型都用一样的随机子.
因此, 要确保$ϵ$在同批次内需生成batch_size个随机数.
所以, 这里生成了一个随机序列: (batch_size $\times 1\times 1\times 1$). 这样就保证每一个速度模型 ($1\times 70\times 70$) 有唯一的随机子.
interpolates和d_interpolates分别对应一个函数的自变量和因变量, 带入到autograd.grad()中就可以返回这个函数相对于自变量的梯度情况. (具体用法可百度) 最终返回的梯度和自变量是同样尺寸大小的 (Same as $\tilde{x}$).
后续对于一张速度模型($1\times 70\times 70$) 的二范数, 应当是直接求每个像素点的平方和再开方. 因此图像本身的维度信息变得无关紧要, 因此代码中通过gradients.view(gradients.size(0), -1)将原本
(batch_size$\times 1\times 70\times 70$)的梯度图像的一张图像给压缩为一维, 最终得到维度为(batch_size$\times 4900$)的压缩后的图像.
后续使用了py中的norm(2, dim=1)方法求二范数, 其含义是对输入矩阵的每行的向量求二范数, 从而将(batch_size$\times 4900$)的压缩后的图像变为(batch_size$\times 1$)的向量, 这个向量中的每个元素都代表了batch中某个速度模型的二范数.

2.5.2 生成器loss

def criterion_g(pred, gt, model_d=None):
    l1loss = nn.L1Loss()
    l2loss = nn.MSELoss()
    loss_g1v = l1loss(pred, gt)
    loss_g2v = l2loss(pred, gt)
    loss = lambda_g1v * loss_g1v + lambda_g2v * loss_g2v
    if model_d is not None:
        loss_adv = -torch.mean(model_d(pred))
        loss += lambda_adv * loss_adv
    return loss, loss_g1v, loss_g2v

torch.mean(model_d(pred))对应了$\underset{\tilde{x}\sim {\mathbb{P}}_g}{\mathbb{E}}D(\tilde{x})$, 他将返回一个数值.
model_d(pred)对应了$D(\tilde{x})$, 因为2022版本抛弃了PatchGAN, 因此这里会返回维度为(batch_size $\times$ 1)的而并非(batch_size $\times$ 4)的向量.
pred对应生成器预测的速度模型$\tilde{x}$维度是(batch_size $\times 1\times 70\times 70)$

2.5.3 训练的配置顺序

for cur_epoch in range(epochs):

    model_g.train()
    model_d.train()

    iter = 1         			# step in current epoch
    max_iter = len(loader)
	...
    start_time = time.time()

    for i, (seismic_data, velocity_models) in enumerate(loader):

        if torch.cuda.is_available():
            seismic_data = seismic_data.cuda(non_blocking=True)
            velocity_models = velocity_models.cuda(non_blocking=True)

        # Fix the generator and update the discriminator first
        optimizer_d.zero_grad()
        with torch.no_grad():	# The generator does not take derivatives
            pred_seismic_data = model_g(seismic_data)
        loss_d, loss_diff, loss_gp = criterion_d(velocity_models, pred_seismic_data, model_d)
        loss_d.backward() 		# Update the discriminator
        optimizer_d.step()

        # Regularly update the generator
        if iter % n_critic == 0 or iter == max_iter:
            optimizer_g.zero_grad()
            pred_seismic_data = model_g(seismic_data)
            loss_g, loss_g1v, loss_g2v = criterion_g(pred_seismic_data, velocity_models, model_d)
            loss_g.backward()
            optimizer_g.step()

        # Exhibition
        ...
		...
        iter += 1
    
    # Statistics and record loss
    ...
    ...
    ...

这里的训练顺序和一般的GAN是类似的.
首先在一轮epoch中套入loader, 再通过loader循环来读取当前epoch中的每个batches. 这个是Pytorch训练的基本方式.
在确定数据资料及其格式正确后, 便先进入判别器的训练.
训练时需要通过 (见下述代码):

optimizer_d.zero_grad()
with torch.no_grad():	# The generator does not take derivatives
	pred_seismic_data = model_g(seismic_data)

来固定生成器G, 即取消生成器G的求导能力. 因为我们担心在针对判别器D求导时因为传递作用而干扰到生成器G. GAN最开始的求导思想是分部执行的.
而后就是训练判别器, 见下述代码↓
在计算loss的时候需要导入判别器本身, 因为loss中有个需要引入网络本身来计算梯度.

loss_d, loss_diff, loss_gp = criterion_d(velocity_models, pred_seismic_data, model_d)
loss_d.backward()     	# Update the discriminator
optimizer_d.step()

当训练的轮次满足n_critic次之后, 就该在当轮额外进行一次生成器loss的计算和反向传播.
这个时候需要通过optimizer_g.zero_grad()初始化生成器的梯度, 即恢复求导能力.
然后进行forward和back propagation.
这个过程无需固定判别器D, 因为我们通过对抗性约束来强化生成器G之余, 其惩罚还会反馈给判别器D来进行强化, 形成对抗.

if iter % n_critic == 0 or iter == max_iter:
	optimizer_g.zero_grad()
	pred_seismic_data = model_g(seismic_data)
	loss_g, loss_g1v, loss_g2v = criterion_g(pred_seismic_data, velocity_models, model_d)
	loss_g.backward()
	optimizer_g.step()

2.6 复现结果

总的来看, VelocityGAN在整体风格上是类似InversionNet的, 但是细节指标上会比InversionNet优越一些 (这里不做展示). 这里我使用了OpenFWI的数据集中的CurveVel A的一部分, 没有用全部的OpenFWI数据 (太多了).

3 存在的主要问题

暂时没有太大的问题.
VelocityGAN的复现效果可能并没有达论文中所示的最佳, 基本猜测可能是数据量的问题.
因为原文中采用了全部的OpenFWI中的数据, 而我暂时只用了部分.
另外原论文还提供了迁移训练, 我们可以通过OpenFWI提供的CurveVel的无断层数据做预训练, 包含断层的CurveFault做迁移训练. 但是因为精力原因并没有尝试去复现.

4 总结及下一步工作

总结

VelocityGAN给人最大的启示其实在于…它其实是基于InversionNet为基础改进得到的, 而且改动不算特别大.
换句话说, 我们似乎可以给一切的端到端网络都添加一个判别器D来构造一个对抗性约束, 从而强化原有网络的学习效率.
这种优化我以前认为可能很复杂, 但是通过接触VelocityGAN后, 我发现这个过程实际上远不如那么麻烦, 因为已有GAN的判别器理论和代码已经非常完美了, 无论是这里的WGAN, WGAN-GP还是其余的GAN结构.

下一步工作

下周将会阅读论文《Multiscale data-driven seismic full-waveform inversion with field data study》
自然图像做FWI, 同时也采用了高频和低频的学习思路, 很有创意的一篇论文.