马坚伟老师基于深度学习对速度反演
FCN-VMB基于深度学习的速度模型反演代码:https://github.com/YangFangShu/FCNVMB-Deep-learning-based-seismic-velocity-model-building
记录学习过程
其中DataLoad_Test.py是加载测试数据集,并对测试数据集进行预处理返回处理后的测试数据和标签
调整数据形状,以匹配模型的输入和输出要求。
块减缩的目的是降低数据维度,使模型处理数据更加高效。
为了验证该方法的稳定性,我们在每个测试地震数据中添加了均值为零、标准导数为5%的高斯噪声。此外,我们将地震数据的振幅提高了两倍。噪声或放大的数据也被用作输入,并被送入网络来反演速度值,用于预测的速度模型都不包含在训练数据集中
# -*- coding: utf-8 -*-
"""
Load testing data set
Created on Feb 2018
@author: fangshuyang (yfs2016@hit.edu.cn)
"""
# 导入所需的库
import numpy as np
from skimage.measure import block_reduce
import skimage
import scipy.io
# 定义DataLoad_Test函数,用于加载和处理测试数据
def DataLoad_Test(test_size,test_data_dir,data_dim,in_channels,model_dim,data_dsp_blk,label_dsp_blk,start,datafilename,dataname,truthfilename,truthname)
'''
#test_size是测试数据集大小在Paramconfig.py里有对应设置;test_data_dir是测试数据集文件地址在#path.py里有;data_dim是炮记录维度(nt,nx);# in_channels=29输入通道数,即放炮数量
in_channels=29 表示每个速度模型真值(即每个标签或地面真实数据)对应29个不同的炮记录数据,意味着每个输入样本由29个不同的地震射线数据(或时间步数据)组成,这些数据将一起用于预测对应的速度模型;
输入和标签:
输入数据:有29个通道的地震数据,每个通道代表不同的地震射线。
标签数据:对应的速度模型真值,是一个二维数据。
ModelDim是速度模型维度(nz,nx) data_dsp_blk:输入的下采样率 label_dsp_blk:输出下采样率;
输入数据:在处理输入数据时,主要关注在时间维度上进行压缩,以减少计算量。这有助于模型更高效地处理大规模数据,同时保留关键的空间特征。
输出数据:在生成输出时,保持与原始数据相同的分辨率,以确保预测结果的高精度和实用性
'''
# 从start开始,循环test_size次
for i in range(start,start+test_size):
# 构建地震数据文件的路径
filename_seis = test_data_dir+'georec_test/'+datafilename+str(i)
#地震数据路径:根目录+'georec_test/'+数据基本文件名+i
print (filename_seis)
# 读取地震数据文件
data1_set = scipy.io.loadmat(filename_seis)
# 将地震数据转换为float32类型,并调整形状
data1_set = np.float32(data1_set[str(dataname)].reshape([data_dim[0],data_dim[1],in_channels]))
'''
把输入的数据变成合适的三维数组 地震数据nt*nx*29
'''
# 对每一个输入通道进行处理和降采样数据处理
for k in range (0,in_channels):
data11_set = np.float32(data1_set[:,:,k]) # 选择第k个通道的数据
data11_set = np.float32(data11_set) # 确保数据类型为float32
# 对数据进行块减缩,使用decimate函数
data11_set = block_reduce(data11_set,block_size=data_dsp_blk,func=decimate)
#对每个通道的地震数据每(5,1)5行1列的数用decimate函数计算得到一个值
data_dsp_dim = data11_set.shape # 获取块减缩后的数据维度
data11_set = data11_set.reshape(1,data_dsp_dim[0]*data_dsp_dim[1]) # 调整数据形状
# 如果是第一个通道,初始化test1_set
if k==0:
test1_set = data11_set
else:
test1_set = np.append(test1_set,data11_set,axis=0) # 追加其他通道的数据
#最终形成(通道数, 高度 * 宽度) 的二维数据集,适合用于神经网络的输入
# 构建标签数据文件的路径
filename_label = test_data_dir+'vmodel_test/'+truthfilename+str(i)
# 读取标签数据文件
data2_set = scipy.io.loadmat(filename_label)
# 将标签数据转换为float32类型,并调整形状
data2_set = np.float32(data2_set[str(truthname)].reshape(model_dim))
# 对标签数据进行块减缩,使用np.max函数
data2_set = block_reduce(data2_set,block_size=label_dsp_blk,func=np.max)
label_dsp_dim = data2_set.shape # 获取块减缩后的标签数据维度
data2_set = data2_set.reshape(1,label_dsp_dim[0]*label_dsp_dim[1]) # 调整标签数据形状
#真值数据,形状为 (1, label_dsp_dim[0] * label_dsp_dim[1])。
data2_set = np.float32(data2_set) # 确保数据类型为float32
# 如果是第一个样本,初始化test_set和label_set
if i==start:
test_set = test1_set
label_set = data2_set
else:
# 追加其他样本的数据
test_set = np.append(test_set,test1_set,axis=0)
label_set = np.append(label_set,data2_set,axis=0)
# 调整测试数据和标签数据的形状
test_set = test_set.reshape((test_size,in_channels,data_dsp_dim[0]*data_dsp_dim[1]))
label_set = label_set.reshape((test_size,label_dsp_dim[0]*label_dsp_dim[1]))
#test_set=(100,29,2000/5*301) label_set=(100,200*301)
# 返回处理后的测试数据、标签数据以及它们的维度
return test_set, label_set, data_dsp_dim, label_dsp_dim
# 定义decimate函数,用于在块减缩时取中间值
def decimate(a,axis):
# 计算每个块的中间值的索引:
idx = np.round((np.array(a.shape)[np.array(axis).reshape(1,-1)]+1.0)/2.0-1).reshape(-1)
# 获取数据的中间值
downa = np.array(a)[:,:,idx[0].astype(int),idx[1].astype(int)]
return downa
DataLoad_Train.py是加载训练数据集,并对训练测试数据集进行预处理,与测试不一样的地方是训练起始是0,测试是start 而且训练集大小与测试集大小以及文件位置名称都不一样 但是整体对数据的处理方式逻辑是一样的
模拟训练数据集包含1600个速度样本,假设每个模型有5到12层作为背景速度,每层的速度值在2000 m/s到4000 m/s之间任意变化。在每个模型中嵌入一个任意形状和位置的盐体,每个盐体的恒定速度值为4500 m/s。每个速度模型的大小采用nx ×n z = 201 × 301网格点,空间间隔dx = dz = 10 m,对于每个速度模型,均匀放置29个源,依次模拟射击集。即每个速度模型对应29个炮记录
UnetModel.py展示了一个用于速度模型构建的U-Net网络的实现,该网络可以从地震非迁移地震数据直接生成速度模型。U-Net是一个经典的卷积神经网络架构,特别适用于图像分割任务。
# -*- coding: utf-8 -*-
"""
Created on Feb 2018
@author: fangshuyang (yfs2016@hit.edu.cn)
"""
################################################
######## DESIGN NETWORK ########
################################################
import torch.nn as nn
import torch
import torch.nn as nn
import torch.nn.functional as F
#定义了一个unetConv2 类是一个包含两个卷积层的网络模块,每个卷积层后可以选择是否包含批量归一化层
class unetConv2(nn.Module):
#实现了两个3x3卷积层,后接批量归一化(BatchNorm)和ReLU激活函数。
def __init__(self, in_size, out_size, is_batchnorm):
'''
#in_size(输入通道数:如果是RGB,则=3), out_size(输出通道数:希望通过该卷积层提取的特征图数量。), is_batchnorm(是否使用批量归一化)
通道数的变化:是特征图深度的变化,而不是图片数量的变化。
特征图深度的变化:由卷积核数量决定,增加通道数可以提取更多特征
'''
super(unetConv2, self).__init__()
#调用 super 的作用是确保父类 nn.Module 的初始化方法被正确调用,从而确保 unetConv2 继承了 #nn.Module 的所有功能和属性
# Kernel size: 3*3, Stride: 1, Padding: 1
#输出特征图的宽度 = (输入特征图宽度 - 卷积核宽度 + 2*填充)/ 步长 + 1 保证经过卷积后图片大小不###变
if is_batchnorm:#判断是否使用批量归一化
self.conv1 = nn.Sequential(nn.Conv2d(in_size, out_size, 3, 1, 1),
nn.BatchNorm2d(out_size),
nn.ReLU(inplace=True),)
self.conv2 = nn.Sequential(nn.Conv2d(out_size, out_size, 3, 1, 1),
nn.BatchNorm2d(out_size),
nn.ReLU(inplace=True),)
else:
self.conv1 = nn.Sequential(nn.Conv2d(in_size, out_size, 3, 1, 1),
nn.ReLU(inplace=True),)
self.conv2 = nn.Sequential(nn.Conv2d(out_size, out_size, 3, 1, 1),
nn.ReLU(inplace=True),)
def forward(self, inputs):
outputs = self.conv1(inputs)
outputs = self.conv2(outputs)
return outputs
#前向传播,经过两层卷积,返回经过两层卷积处理后的输出
class unetDown(nn.Module):
def __init__(self, in_size, out_size, is_batchnorm):
super(unetDown, self).__init__()
self.conv = unetConv2(in_size, out_size, is_batchnorm)
self.down = nn.MaxPool2d(2, 2, ceil_mode=True)
def forward(self, inputs):
outputs = self.conv(inputs)
outputs = self.down(outputs)
return outputs
'''
unetDown 类是 U-Net 网络的一个下采样模块,包含一个卷积层,包括两次卷积
一个最大池化层。
它先通过卷积层提取特征,再通过最大池化层下采样,从而减少特征图的尺寸,聚合特征,提高模型的感受野。
'''
class unetUp(nn.Module):
#定义一个上采样模块,用于对特征图进行上采样和融合
def __init__(self, in_size, out_size, is_deconv):
super(unetUp, self).__init__()
self.conv = unetConv2(in_size, out_size, True)
#通过两个卷积层实现的卷积操作
# Transposed convolution
#self.up:上采样操作,若 is_deconv 为 True,使用转置卷积;否则使用双线性上采样。
if is_deconv:
self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2,stride=2)
#nn.ConvTranspose2d 是 PyTorch 中的一个模块,用于实现转置卷积(也称为反卷积或上卷积)。转置卷积#通常用于增加特征图的空间分辨率,kernel_size=2: 卷积核大小为 2x2。stride=2: 步幅为 2,这意味着#输出特征图的空间尺寸是输入特征图的两倍。同时减少通道数,以便与下采样路径中的特征图进行拼接。
else:
self.up = nn.UpsamplingBilinear2d(scale_factor=2)
#nn.UpsamplingBilinear2d 是 PyTorch 中的一种上采样层,使用双线性插值法来增加特征图的空间分辨###率。这种方法不会引入额外的可学习参数,与卷积和转置卷积不同,插值方法通过平滑地插入新像素值来实#现上采样
#scale_factor=2: 表示输出特征图的宽
'''
forward 方法通过上采样、填充、拼接和卷积的组合操作,实现了 U-Net 网络中从编码器到解码器的特征图传递和融合。度和高度是输入特征图的两倍
'''
def forward(self, inputs1, inputs2):
#在 U-Net 中,这些输入通常是来自编码器的特征图(inputs1)和来自上一层解码器的特征图(inputs2)
outputs2 = self.up(inputs2)
#对inputs2进行上采样大小乘2
offset1 = (outputs2.size()[2]-inputs1.size()[2])#计算高度差
offset2 = (outputs2.size()[3]-inputs1.size()[3])#计算宽度差
padding=[offset2//2,(offset2+1)//2,offset1//2,(offset1+1)//2]#看需要补多少边以#inputs1与outputs2大小相等
# Skip and concatenate
outputs1 = F.pad(inputs1, padding)#对 inputs1进行padding使其与 outputs2 尺寸匹配。
return self.conv(torch.cat([outputs1, outputs2], 1))
#将经过 padding 的 inputs1 和上采样后的 inputs2 在通道维度上拼接。拼接后结果经过卷积层处理,以提#取融合后的特征。
'''
编码器(Encoder)
特征提取:
编码器由多个卷积层(通常配合激活函数如 ReLU)和池化层组成。
每个卷积层负责提取输入图像的特征,而池化层负责降低特征图的空间分辨率(即图像的宽度和高度)。
通过这些操作,编码器将输入图像逐渐转换为一系列逐渐降低分辨率但包含丰富特征信息的特征图。
解码器(Decoder)
特征恢复:
解码器的任务是逐步恢复图像的空间分辨率,最终恢复到与输入图像相同的尺寸。
通过使用上采样层(如反卷积或双线性插值)来增加特征图的空间分辨率
'''
class UnetModel(nn.Module):
#初始化参数UnetModel 类的初始化
def __init__(self, n_classes, in_channels ,is_deconv, is_batchnorm):
super(UnetModel, self).__init__()
self.is_deconv = is_deconv#是否使用反卷积
self.in_channels = in_channels#输入通道数
self.is_batchnorm = is_batchnorm#是否使用批量归一化
self.n_classes =n_classes
filters = [64, 128, 256, 512, 1024]
'''
这个列表定义了 U-Net 网络中每个卷积层的输出通道数(即特征图的深度,卷积核个数),从而控制每层卷积的输出特征图的数量。
通常,随着网络深度的增加,特征图的数量(通道数)会增加,以便捕捉更多的高层特征信息
'''
#unetDown 是编码器模块,用于下采样输入数据,提取特征并减少空间维度。
self.down1 = unetDown(self.in_channels, filters[0], self.is_batchnorm)
#输入通道数:29 输出:64 尺寸缩小一半 400*301-200*151
self.down2 = unetDown(filters[0], filters[1], self.is_batchnorm)
#输入通道数:64 输出:128 尺寸缩小一半 100*75
self.down3 = unetDown(filters[1], filters[2], self.is_batchnorm)
#输入通道数:128 输出:256 尺寸缩小一半 50*38
self.down4 = unetDown(filters[2], filters[3], self.is_batchnorm)
#输入通道数:256 输出:512 尺寸缩小一半 25*19
self.center = unetConv2(filters[3], filters[4], self.is_batchnorm)
#输入通道数:512 输出:1024 尺寸不变 25*19
#处理编码器中最后一层的输出,为解码器部分准备特征图。
#unetUp 是解码器模块,用于上采样特征图并恢复空间分辨率。
self.up4 = unetUp(filters[4], filters[3], self.is_deconv)
#输入通道数:1024 输出:512 尺寸扩大一倍 50*38
self.up3 = unetUp(filters[3], filters[2], self.is_deconv)
#输入通道数:512 输出:256 尺寸扩大一倍 100*76
self.up2 = unetUp(filters[2], filters[1], self.is_deconv)
#输入通道数:256 输出:128 尺寸扩大一倍 200*152
self.up1 = unetUp(filters[1], filters[0], self.is_deconv)
#输入通道数:128 输出:64 尺寸扩大一倍 400*304
self.final = nn.Conv2d(filters[0],self.n_classes, 1)
'''
在这个针对地震速度模型构建(VMB)修改的U-Net中,`Nclasses = 1` 的解释如下:
1. **多通道输入**:标准U-Net处理的是RGB图像,每个图像有三个颜色通道。在这个修改后的U-Net中,输入的每个通道代表不同的地震炮记录(shot gathers),这些记录来自相同的模型,因此输入通道的数量由地震炮记录的数量决定=29。
2. **单输出通道**:修改后的U-Net的输出旨在表示一个速度模型。由于速度模型是一个单通道的图像(每个空间位置只有一个值),因此输出层的通道数设置为 `Nclasses = 1`。这意味着网络将输出一个单通道的图像,直接对应于每个位置的预测速度值。
3. **领域投影**:该网络旨在将地震数据从 (x, t) 域(空间-时间域)投影到 (x, z) 域(空间-深度域),有效地将数据转换并映射到所需的速度模型。
4. **特征图尺寸调整**:网络在最终卷积层后的特征图尺寸调整为与速度模型的尺寸相匹配,确保输出与预测的速度模型大小一致。
通过将 `Nclasses = 1` 设置为1,网络配置为输出每个速度模型像素的单一值,这与预测连续的单通道输出的目标一致。
'''
def forward(self, inputs,label_dsp_dim):
down1 = self.down1(inputs)
down2 = self.down2(down1)
down3 = self.down3(down2)
down4 = self.down4(down3)
center = self.center(down4)
up4 = self.up4(down4, center)
up3 = self.up3(down3, up4)
up2 = self.up2(down2, up3)
up1 = self.up1(down1, up2)
up1 = up1[:,:,1:1+label_dsp_dim[0],1:1+label_dsp_dim[1]].contiguous()
#将上采样后的特征图裁剪到目标尺寸(label_dsp_dim),以确保输出的特征图尺寸与实际需要的标签尺寸一#致。这里进行裁剪是为了匹配最终预测的标签尺寸,避免由于卷积和上采样操作导致的边界效应。
return self.final(up1)
#self.final(up1):通过最后一个卷积层,将特征图映射到输出通道的数量(self.n_classes),生成最终#的预测结果。在这个特定的应用中,self.n_classes 为1,表示单通道的速度模型输出
# Initialization of Parameters
def _initialize_weights(self):
for m in self.modules():
#遍历所有模块
if isinstance(m, nn.Conv2d):
#是否是卷积模块
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
'''
计算卷积核的总参数数量,用于初始化权重。这里的 m.kernel_size 是卷积核的大小(宽和高),m.out_channels 是输出通道数。
'''
m.weight.data.normal_(0, sqrt(2. / n))
#使用正态分布初始化卷积层的权重。均值为0,标准差为 sqrt(2. / n)。
if m.bias is not None:
m.bias.data.zero_()
#检查是否有偏置项。将偏置初始化为0
elif isinstance(m, nn.BatchNorm2d):
#检查是否为批量归一化层
m.weight.data.fill_(1)
#将批量归一化层的权重初始化为1。
m.bias.data.zero_()
#将批量归一化层的偏置初始化为0。
elif isinstance(m,nn.ConvTranspose2d):
#检查当前模块是否是转置卷积层。
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
'''
计算卷积核的总参数数量,用于初始化权重。这里的 m.kernel_size 是卷积核的大小(宽和高),m.out_channels 是输出通道数。
'''
m.weight.data.normal_(0, sqrt(2. / n))
#使用正态分布初始化卷积层的权重。均值为0,标准差为 sqrt(2. / n)。
if m.bias is not None:
m.bias.data.zero_()
#检查是否有偏置项。将偏置初始化为0
utils.py:这段代码包含了用于计算和保存模型评估指标、处理图像、以及可视化结果的功能。下面是各个部分的详细解释
# -*- coding: utf-8 -*-
"""
Created on Feb 2018
@author: fangshuyang (yfs2016@hit.edu.cn)
"""
import torch
import numpy as np
import torch.nn as nn
from math import log10
from torch.autograd import Variable
from math import exp
import torch.nn.functional as F
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import scipy.io
from mpl_toolkits.axes_grid1 import make_axes_locatable
def turn(GT):
dim = GT.shape
for j in range(0,dim[1]):
for i in range(0,dim[0]//2):
temp = GT[i,j]
GT[i,j] = GT[dim[0]-1-i,j]
GT[dim[0]-1-i,j] = temp
return GT
#翻转图像的函数。
def PSNR(prediction, target):
prediction = Variable(torch.from_numpy(prediction))
target = Variable(torch.from_numpy(target))
zero = torch.zeros_like(target)
criterion = nn.MSELoss(size_average=True)
MSE = criterion (prediction, target)
total = criterion (target, zero)
psnr = 10. * log10(total.item() / MSE.item())
return psnr
def gaussian(window_size, sigma):
gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
return gauss / gauss.sum()
def create_window(window_size, channel):
_1D_window = gaussian(window_size, 1.5).unsqueeze(1)
_2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
return window
def _ssim(img1, img2, window, window_size, channel, size_average=True):
mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
mu1_sq = mu1.pow(2)
mu2_sq = mu2.pow(2)
mu1_mu2 = mu1 * mu2
sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2
L = 255
C1 = (0.01*L) ** 2
C2 = (0.03*L) ** 2
ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
if size_average:
return ssim_map.mean()
else:
return ssim_map.mean(1).mean(1).mean(1)
'''
SSIM 地图:ssim_map 是一张与输入图像大小相同的矩阵,其中每个位置的值表示该位置局部区域的结构相似性。这种地图可以显示图像中不同区域的相似性强度,帮助分析哪些区域的质量较好或较差
'''
#结果相似性
def SSIM(img1, img2, window_size=11, size_average=True):
img1 = Variable(torch.from_numpy(img1))
img2 = Variable(torch.from_numpy(img2))
(_, channel, _, _) = img1.size()
window = create_window(window_size, channel)
if img1.is_cuda:
window = window.cuda(img1.get_device())
window = window.type_as(img1)
return _ssim(img1, img2, window, window_size, channel, size_average)
'''
函数 SaveTrainResults 的目的是将训练过程中的损失(loss)曲线保存为图像文件,并将损失数据保存为 .mat 文件。以下是该函数的详细解释
'''
def SaveTrainResults(loss,SavePath,font2,font3):
fig,ax = plt.subplots()
plt.plot(loss[1:], linewidth=2)
ax.set_xlabel('Num. of epochs', font2)
ax.set_ylabel('MSE Loss', font2)
ax.set_title('Training', font3)
ax.set_xlim([1,6])
ax.set_xticklabels(('0','20','40','60','80','100'))
#设置 x 轴的显示范围为 1 到 6,并且设置 x 轴刻度标签。
for label in ax.get_xticklabels()+ax.get_yticklabels():
label.set_fontsize(12)
ax.grid(linestyle='dashed',linewidth=0.5)
#保存路径
plt.savefig(SavePath+'TrainLoss', transparent = True)
data = {}
data['loss'] = loss
scipy.io.savemat(SavePath+'TrainLoss',data)
plt.show(fig)
plt.close()
#将损失数据保存为 MATLAB .mat 文件,文件名为 TrainLoss
'''
SaveTestResults 函数的主要作用是保存模型的测试结果,包括总的 PSNR 和 SSIM 值、真实值(Ground Truth)以及预测值。这些结果以 .mat 格式保存在指定的路径中
'''
def SaveTestResults(TotPSNR,TotSSIM,Prediction,GT,SavePath):
'''
TotPSNR:一个包含所有测试样本 PSNR 值的数组。
TotSSIM:一个包含所有测试样本 SSIM 值的数组。
Prediction:模型的预测结果数组。
GT:真实标签数组(Ground Truth)。
SavePath:保存结果文件的路径
'''
data = {}
data['TotPSNR'] = TotPSNR
data['TotSSIM'] = TotSSIM
data['GT'] = GT
data['Prediction'] = Prediction
scipy.io.savemat(SavePath+'TestResults',data)
'''
PlotComparison 函数的主要作用是绘制并保存模型预测结果和真实值(Ground Truth)的对比图。该函数接受模型预测值和真实值,并将它们绘制成图像保存下来,以便进行视觉对比
'''
def PlotComparison(pd,gt,label_dsp_dim,label_dsp_blk,dh,minvalue,maxvalue,font2,font3,SavePath):
'''
pd:模型的预测结果。
gt:真实标签(Ground Truth)。
label_dsp_dim:标签的显示维度,包含两个值 [height, width]。
label_dsp_blk:标签的显示块大小,包含两个值 [block_height, block_width]。
dh:深度高度的比例因子,用于将像素转换为实际单位。
minvalue 和 maxvalue:图像的最小值和最大值,用于设置颜色条的范围。
font2 和 font3:用于设置图像中标签和标题的字体属性。
SavePath:保存图像的路径。
'''
PD = pd.reshape(label_dsp_dim[0],label_dsp_dim[1])
GT = gt.reshape(label_dsp_dim[0],label_dsp_dim[1])
fig1,ax1 = plt.subplots(figsize=(6, 4))
im1 = ax1.imshow(GT,extent=[0,label_dsp_dim[1]*label_dsp_blk[1]*dh/1000., \
# 绘制 Ground Truth 图像 0,label_dsp_dim[0]*label_dsp_blk[0]*dh/1000.],vmin=minvalue,vmax=maxvalue)
divider = make_axes_locatable(ax1)
cax1 = divider.append_axes("right",size="5%",pad=0.05)
plt.colorbar(im1,ax=ax1,cax=cax1).set_label('Velocity (m/s)')
plt.tick_params(labelsize=12)
for label in ax1.get_xticklabels()+ax1.get_yticklabels():
label.set_fontsize(14)
ax1.set_xlabel('Position (km)',font2)
ax1.set_ylabel('Depth (km)',font2)
ax1.set_title('Ground truth',font3)
ax1.invert_yaxis()
plt.subplots_adjust(bottom=0.15,top=0.92,left=0.08,right=0.98)
plt.savefig(SavePath+'GT',transparent=True)
fig2,ax2=plt.subplots(figsize=(6, 4))
im2=ax2.imshow(PD,extent=[0,label_dsp_dim[1]*label_dsp_blk[1]*dh/1000., \
# 绘制 预测 图像 0,label_dsp_dim[0]*label_dsp_blk[0]*dh/1000.],vmin=minvalue,vmax=maxvalue)
plt.tick_params(labelsize=12)
for label in ax2.get_xticklabels()+ax2.get_yticklabels():
label.set_fontsize(14)
ax2.set_xlabel('Position (km)',font2)
ax2.set_ylabel('Depth (km)',font2)
ax2.set_title('Prediction',font3)
ax2.invert_yaxis()
plt.subplots_adjust(bottom=0.15,top=0.92,left=0.08,right=0.98)
plt.savefig(SavePath+'PD',transparent=True)
plt.show(fig1)
plt.show(fig2)
plt.close()
#执行此代码后,预测结果和真实值的对比图将被保存为 ./results/GT.png 和 ./results/PD.png 文件
FCNVMB_train.py:描述了一个用于从预堆叠未迁移地震数据构建速度模型的全卷积神经网络(U-Net)的完整训练过程
# -*- coding: utf-8 -*-
"""
Fully Convolutional neural network (U-Net) for velocity model building from prestack
unmigrated seismic data directly
Created on Feb 2018
@author: fangshuyang (yfs2016@hit.edu.cn)
"""
################################################
######## IMPORT LIBARIES ########
################################################
'''
导入库:从 ParamConfig、PathConfig 和 LibConfig 中导入必要的配置和库。
'''
from ParamConfig import *
from PathConfig import *
from LibConfig import *
################################################
######## NETWORK ########
################################################
# Here indicating the GPU you want to use. if you don't have GPU, just leave it.
##检查 GPU 是否可用,并选择设备。
cuda_available = torch.cuda.is_available()
device = torch.device("cuda" if cuda_available else "cpu")
#初始化 U-Net 模型并将其移动到设备(CPU 或 GPU)。
net = UnetModel(n_classes=Nclasses,in_channels=Inchannels,is_deconv=True,is_batchnorm=True)
if torch.cuda.is_available():
net.cuda()
#定义 Adam 优化器。
# Optimizer we want to use
optimizer = torch.optim.Adam(net.parameters(),lr=LearnRate)
#如果需要重用之前训练的模型,则加载预训练模型。
'''
加载预训练模型的过程实际上就是将已经训练好的参数(权重和偏置)加载到你设置的网络中。这些参数通常保存在一个文件中(如 .pkl 或 .pt 文件),在你需要继续训练、进行迁移学习或推理时,将这些参数载入模型可以省去重新训练的时间和计算资源。
'''
# If ReUse, it will load saved model from premodelfilepath and continue to train
if ReUse:
print('***************** Loading the pre-trained model *****************')
print('')
premodel_file = models_dir + premodelname + '.pkl'
##Load generator parameters
net = net.load_state_dict(torch.load(premodel_file))
net = net.to(device)
print('Finish downloading:',str(premodel_file))
################################################
######## LOADING TRAINING DATA ########
################################################
print('***************** Loading Training DataSet *****************')
#加载训练的数据集
train_set,label_set,data_dsp_dim,label_dsp_dim = DataLoad_Train(train_size=TrainSize,train_data_dir=train_data_dir, \
data_dim=DataDim,in_channels=Inchannels, \
model_dim=ModelDim,data_dsp_blk=data_dsp_blk, \
label_dsp_blk=label_dsp_blk,start=1, \
datafilename=datafilename,dataname=dataname, \
truthfilename=truthfilename,truthname=truthname)
'''
DataLoad_Train 是一个自定义的数据加载函数,用于从指定目录读取训练数据和标签。它会处理数据,并返回以下内容:
train_set:训练数据集
label_set:标签数据集
data_dsp_dim:数据的尺寸
label_dsp_dim:标签的尺寸
输入参数:
train_size:训练数据的数量。
train_data_dir:训练数据的存储目录。
data_dim:原始数据的维度。
in_channels:输入通道的数量。
model_dim:模型标签的维度。
data_dsp_blk:数据块的尺寸。
label_dsp_blk:标签块的尺寸。
start:数据加载的起始位置(通常是 1)。
datafilename 和 truthfilename:包含数据和标签文件的文件名。
dataname 和 truthname:数据和标签的具体名称
'''
# Change data type (numpy --> tensor)
train = data_utils.TensorDataset(torch.from_numpy(train_set),torch.from_numpy(label_set))
train_loader = data_utils.DataLoader(train,batch_size=BatchSize,shuffle=True)
#data_utils.TensorDataset:将数据和标签封装到一个 TensorDataset 对象中。这是 PyTorch 提供的一#个类,它将数据和标签配对,使得每次迭代都可以获取数据和对应的标签
################################################
######## TRAINING ########
################################################
print()
print('*******************************************')
print('*******************************************')
print(' START TRAINING ')
print('*******************************************')
print('*******************************************')
print()
'''
开始训练
打印训练参数(数据维度、训练大小、批量大小、学习率等)。
迭代训练数据,执行前向传播、计算损失、反向传播并更新模型参数。
每隔 DisplayStep 步打印一次损失值。
每完成一个 epoch 后,打印该 epoch 的损失值和耗时。
每隔 SaveEpoch 步保存一次模型参数
'''
print ('Original data dimention:%s' % str(DataDim))
print ('Downsampled data dimention:%s ' % str(data_dsp_dim))
print ('Original label dimention:%s' % str(ModelDim))
print ('Downsampled label dimention:%s' % str(label_dsp_dim))
print ('Training size:%d' % int(TrainSize))
print ('Traning batch size:%d' % int(BatchSize))
print ('Number of epochs:%d' % int(Epochs))
print ('Learning rate:%.5f' % float(LearnRate))
# Initialization
loss1 = 0.0
'''
用于累积每个 epoch 的平均损失。在训练过程中,它会被更新以记录和保存损失值。这通常用于记录和可视化训练过程中的损失变化情况
'''
step = np.int(TrainSize/BatchSize)
#可以得出每个 epoch 中需要多少次迭代来处理所有训练样本
start = time.time()
#用于后续计算训练过程所需的时间
for epoch in range(Epochs):
epoch_loss = 0.0#计算每个epoch对应的损失函数
since = time.time()
#遍历训练数据加载器 train_loader 中的批次。images 和 labels 分别是当前批次的输入数据和标签。
for i, (images,labels) in enumerate(train_loader):
iteration = epoch*step+i+1
# Set Net with train condition
net.train()
#将模型设置为训练模式
# Reshape data size
images = images.view(BatchSize,Inchannels,data_dsp_dim[0],data_dsp_dim[1])
labels = labels.view(BatchSize,Nclasses,label_dsp_dim[0],label_dsp_dim[1])
images = images.to(device)
labels = labels.to(device)
# Zero the gradient buffer
optimizer.zero_grad()
# Forward prediction
'''
outputs = net(images, label_dsp_dim): 将 images 输入到模型 net 中进行前向传播,得到预测结果 outputs。label_dsp_dim 是模型的输出尺寸要求,传递给模型以确保输出尺寸的匹配
'''
outputs = net(images,label_dsp_dim)
'''
这段代码计算了模型的损失,进行反向传播并优化模型,最后在指定的迭代步数上打印训练损失。逐行解释如下
'''
# Calculate the MSE
loss = F.mse_loss(outputs,labels,reduction='sum')/(label_dsp_dim[0]*label_dsp_dim[1]*BatchSize)
#计算均方误差 (MSE) 损失,其中 outputs 是模型的预测结果,labels 是真实标签。以获得每个样本##的平均损失。这是为了确保损失值的缩放是合理的,尤其是当数据集的规模发生变化时。
if np.isnan(float(loss.item())):
raise ValueError('loss is nan while training')
#: 检查计算得到的损失是否为 NaN (不是一个数字),这可能是由于数值不稳定或其他计算问题导致的
epoch_loss += loss.item()
#将当前批次的损失值加到 epoch_loss 中,以便在每个 epoch 完成时计算总损失
# Loss backward propagation
loss.backward()
# 方法会根据损失值计算各个模型参数的梯度
# Optimize
optimizer.step()
#optimizer.step():更新模型的参数。这个步骤使用之前计算的梯度来调整模型的权重,以减少损失函数的值
# Print loss
if iteration % DisplayStep == 0:
#DisplayStep = 2: 这表示每经过两个迭代步骤,代码会输出一次训练损失和其他相关统计信息
print('Epoch: {}/{}, Iteration: {}/{} --- Training Loss:{:.6f}'.format(epoch+1, \
Epochs,iteration, \
step*Epochs,loss.item()))
#每隔一定的迭代步数(DisplayStep),打印一次损失值。
#打印当前 epoch、迭代次数以及训练损失。{:.6f} 指定以浮点数格式打印损失值,保留六位小数。
# Print loss and consuming time every epoch
if (epoch+1) % 1 == 0:
#print ('Epoch [%d/%d], Loss: %.10f' % (epoch+1,Epochs,loss.item()))
#loss1 = np.append(loss1,loss.item())
print('Epoch: {:d} finished ! Loss: {:.5f}'.format(epoch+1,epoch_loss/i))
loss1 = np.append(loss1,epoch_loss/i)
time_elapsed = time.time() - since
print('Epoch consuming time: {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
#定期保存模型的参数
# Save net parameters every 10 epochs
if (epoch+1) % SaveEpoch == 0:
torch.save(net.state_dict(),models_dir+modelname+'_epoch'+str(epoch+1)+'.pkl')
print ('Trained model saved: %d percent completed'% int((epoch+1)*100/Epochs))
#这段代码的功能是记录训练的总时间,并保存训练过程中记录的损失值。具体步骤如下:
# Record the consuming time
time_elapsed = time.time() - start
print('Training complete in {:.0f}m {:.0f}s' .format(time_elapsed //60 , time_elapsed % 60))
# Save the loss
font2 = {'family': 'Times New Roman',
'weight': 'normal',
'size': 17,
}
font3 = {'family': 'Times New Roman',
'weight': 'normal',
'size': 21,
}
SaveTrainResults(loss=loss1,SavePath=results_dir,font2=font2,font3=font3)
# 函数的作用是将训练过程中记录的损失值以图像和数据文件的形式保存到指定的路径
FCN_VMB_test:测试阶段,加载预训练的U-Net模型,并对测试数据进行测试,计算并保存预测结果。
# -*- coding: utf-8 -*-
"""
Fully Convolutional neural network (U-Net) for velocity model building from prestack
unmigrated seismic data directly
Created on Feb 2018
@author: fangshuyang (yfs2016@hit.edu.cn)
"""
################################################
######## IMPORT LIBARIES ########
################################################
from ParamConfig import *
from PathConfig import *
from LibConfig import *
################################################
######## LOAD NETWORK ########
################################################
# Here indicating the GPU you want to use. if you don't have GPU, just leave it.
cuda_available = torch.cuda.is_available()
device = torch.device("cuda" if cuda_available else "cpu")
#模型文件: 定义了模型文件的路径,该文件是保存了训练好的模型参数的。
model_file = models_dir+modelname+'_epoch'+str(Epochs)+'.pkl'
net = UnetModel(n_classes=Nclasses,in_channels=Inchannels, \
is_deconv=True,is_batchnorm=True)
#加载模型: 实例化U-Net模型并加载保存的参数
net.load_state_dict(torch.load(model_file))
if torch.cuda.is_available():
net.cuda()
################################################
######## LOADING TESTING DATA ########
################################################
print('***************** Loading Testing DataSet *****************')
'''
加载测试数据: 从指定的数据目录中加载测试数据集,并将其转换为PyTorch的TensorDataset和DataLoader,便于批量处理
'''
test_set,label_set,data_dsp_dim,label_dsp_dim = DataLoad_Test(test_size=TestSize,test_data_dir=test_data_dir, \
data_dim=DataDim,in_channels=Inchannels, \
model_dim=ModelDim,data_dsp_blk=data_dsp_blk, \
label_dsp_blk=label_dsp_blk,start=1601, \
datafilename=datafilename,dataname=dataname, \
truthfilename=truthfilename,truthname=truthname)
test = data_utils.TensorDataset(torch.from_numpy(test_set),torch.from_numpy(label_set))
test_loader = data_utils.DataLoader(test,batch_size=TestBatchSize,shuffle=False)
################################################
######## TESTING ########
################################################
print()
print('*******************************************')
print('*******************************************')
print(' START TESTING ')
print('*******************************************')
print('*******************************************')
print()
# Initialization
#初始化: 计算总测试时间,初始化用于存储PSNR、SSIM、预测结果和真实标签的数组。
since = time.time()
TotPSNR = np.zeros((1,TestSize),dtype=float)
TotSSIM = np.zeros((1,TestSize),dtype=float)
Prediction = np.zeros((TestSize,label_dsp_dim[0],label_dsp_dim[1]),dtype=float)
GT = np.zeros((TestSize,label_dsp_dim[0],label_dsp_dim[1]),dtype=float)
total = 0
#测试循环: 遍历测试数据集,通过模型进行预测,并计算每个样本的PSNR和SSIM指标
for i, (images,labels) in enumerate(test_loader):
images = images.view(TestBatchSize,Inchannels,data_dsp_dim[0],data_dsp_dim[1])
labels = labels.view(TestBatchSize,Nclasses,label_dsp_dim[0],label_dsp_dim[1])
images = images.to(device)
labels = labels.to(device)
#保存结果: 将计算得到的PSNR、SSIM、预测结果和真实标签保存到文件中
# Predictions
net.eval()
outputs = net(images,label_dsp_dim)
outputs = outputs.view(TestBatchSize,label_dsp_dim[0],label_dsp_dim[1])
outputs = outputs.data.cpu().numpy()
gts = labels.data.cpu().numpy()
# Calculate the PSNR, SSIM
for k in range(TestBatchSize):
pd = outputs[k,:,:].reshape(label_dsp_dim[0],label_dsp_dim[1])
gt = gts[k,:,:].reshape(label_dsp_dim[0],label_dsp_dim[1])
pd = turn(pd)
gt = turn(gt)
Prediction[i*TestBatchSize+k,:,:] = pd
GT[i*TestBatchSize+k,:,:] = gt
psnr = PSNR(pd,gt)
TotPSNR[0,total] = psnr
ssim = SSIM(pd.reshape(-1,1,label_dsp_dim[0],label_dsp_dim[1]),gt.reshape(-1,1,label_dsp_dim[0],label_dsp_dim[1]))
TotSSIM[0,total] = ssim
print('The %d testing psnr: %.2f, SSIM: %.4f ' % (total,psnr,ssim))
total = total + 1
# Save Results
SaveTestResults(TotPSNR,TotSSIM,Prediction,GT,results_dir)
#绘图: 绘制预测结果与真实标签的比较图。
# Plot one prediction and ground truth
num = 0
if SimulateData:
minvalue = 2000
else:
minvalue = 1500
maxvalue = 4500
font2 = {'family': 'Times New Roman',
'weight': 'normal',
'size': 17,
}
font3 = {'family': 'Times New Roman',
'weight': 'normal',
'size': 21,
}
PlotComparison(Prediction[num,:,:],GT[num,:,:],label_dsp_dim,label_dsp_blk,dh,minvalue,maxvalue,font2,font3,SavePath=results_dir)
#记录时间: 输出测试过程所用的时间
# Record the consuming time
time_elapsed = time.time() - since
print('Testing complete in {:.0f}m {:.0f}s' .format(time_elapsed // 60, time_elapsed % 60))