当前位置：首页 > article >正文

介绍使用 WGAN（Wasserstein GAN）网络对天然和爆破的地震波形图进行分类的实现步骤

article 2025/2/2 6:50:14

以下将为你详细介绍使用 WGAN（Wasserstein GAN）网络对天然和爆破的地震波形图进行分类的实现步骤，包含代码实现和项目结题报告的大纲。

代码实现

1. 环境准备

确保你已经安装了必要的库，如 torch、torchvision、numpy、matplotlib 等。可以使用以下命令安装：

pip install torch torchvision numpy matplotlib

2. 代码示例

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt

# 自定义地震波形数据集类
class SeismicDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data[idx]
        label = self.labels[idx]
        return torch.tensor(sample, dtype=torch.float32), torch.tensor(label, dtype=torch.long)

# 生成器网络
class Generator(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.LeakyReLU(0.2),
            nn.Linear(128, 256),
            nn.BatchNorm1d(256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 512),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, output_dim),
            nn.Tanh()
        )

    def forward(self, z):
        return self.model(z)

# 判别器网络
class Discriminator(nn.Module):
    def __init__(self, input_dim):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 1)
        )

    def forward(self, x):
        return self.model(x)

# 训练函数
def train_wgan(generator, discriminator, dataloader, num_epochs, lr, lambda_gp, batch_size, z_dim):
    gen_optimizer = optim.RMSprop(generator.parameters(), lr=lr)
    disc_optimizer = optim.RMSprop(discriminator.parameters(), lr=lr)

    for epoch in range(num_epochs):
        for i, (real_data, _) in enumerate(dataloader):
            real_data = real_data.view(real_data.size(0), -1)

            # 训练判别器
            for _ in range(5):
                z = torch.randn(batch_size, z_dim)
                fake_data = generator(z)

                disc_optimizer.zero_grad()
                real_output = discriminator(real_data)
                fake_output = discriminator(fake_data.detach())

                # 计算 Wasserstein 距离
                wasserstein_distance = real_output.mean() - fake_output.mean()

                # 计算梯度惩罚
                alpha = torch.rand(real_data.size(0), 1)
                alpha = alpha.expand(real_data.size())
                interpolates = alpha * real_data + ((1 - alpha) * fake_data.detach())
                interpolates.requires_grad = True
                interpolates_output = discriminator(interpolates)

                gradients = torch.autograd.grad(
                    outputs=interpolates_output,
                    inputs=interpolates,
                    grad_outputs=torch.ones(interpolates_output.size()),
                    create_graph=True,
                    retain_graph=True,
                    only_inputs=True
                )[0]
                gradients = gradients.view(gradients.size(0), -1)
                gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()

                disc_loss = -wasserstein_distance + lambda_gp * gradient_penalty
                disc_loss.backward()
                disc_optimizer.step()

            # 训练生成器
            gen_optimizer.zero_grad()
            z = torch.randn(batch_size, z_dim)
            fake_data = generator(z)
            fake_output = discriminator(fake_data)
            gen_loss = -fake_output.mean()
            gen_loss.backward()
            gen_optimizer.step()

        print(f'Epoch [{epoch + 1}/{num_epochs}] Gen Loss: {gen_loss.item():.4f} Disc Loss: {disc_loss.item():.4f}')

# 分类函数
def classify_seismic_waves(generator, discriminator, test_data):
    test_data = test_data.view(test_data.size(0), -1)
    outputs = discriminator(test_data)
    predictions = (outputs > 0).float()
    return predictions

# 示例数据（需要替换为真实的地震波形数据）
num_samples = 1000
waveform_length = 100
data = np.random.randn(num_samples, waveform_length)
labels = np.random.randint(0, 2, num_samples)

# 创建数据集和数据加载器
dataset = SeismicDataset(data, labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 初始化网络
z_dim = 100
generator = Generator(z_dim, waveform_length)
discriminator = Discriminator(waveform_length)

# 训练 WGAN
train_wgan(generator, discriminator, dataloader, num_epochs=10, lr=0.00005, lambda_gp=10, batch_size=32, z_dim=z_dim)

# 分类测试数据（这里简单使用训练数据作为测试数据示例）
test_data = torch.tensor(data[:10], dtype=torch.float32)
predictions = classify_seismic_waves(generator, discriminator, test_data)
print("分类结果:", predictions)