pytorch实现gan扩充莺尾花数据集代码

以下是使用 PyTorch 实现 GAN 对鸢尾花数据集进行扩充的示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np


class Generator(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        
        self.gen = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, output_dim),
            nn.Tanh()
        )

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


class Discriminator(nn.Module):
    def __init__(self, input_dim):
        super().__init__()

        self.disc = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )

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


# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Hyper-parameters
batch_size = 64
input_dim_g = 100  # Input noise dimension for generator
input_dim_d = 4    # Input data dimension for discriminator (iris dataset has 4 features)
output_dim_g = 4   # Output data dimension for generator (iris dataset has 4 features)
lr = 0.0002
num_epochs = 200


# Load the iris dataset
def load_data():
    transform = transforms.Compose([
        transforms.ToTensor(),
    ])
    train_dataset = datasets.load_iris(root="./data", train=True, download=True, transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    return train_loader


def train(generator, discriminator, train_loader):
    # Loss functions and optimizers
    criterion = nn.BCELoss()
    optimizer_g = optim.Adam(generator.parameters(), lr=lr)
    optimizer_d = optim.Adam(discriminator.parameters(), lr=lr)

    for epoch in range(num_epochs):
        for batch_idx, (real_data, _) in enumerate(train_loader):
            real_data = real_data.view(-1, 4).to(device)

            # Train discriminator: max log(D(x)) + log(1 - D(G(z)))
            noise = torch.randn(batch_size, input_dim_g).to(device)
            fake_data = generator(noise)
            label_real = torch.ones(batch_size, 1).to(device)
            label_fake = torch.zeros(batch_size, 1).to(device)

            # Forward pass real and fake data through discriminator separately
            output_real = discriminator(real_data)
            output_fake = discriminator(fake_data)

            # Calculate the loss for discriminator
            loss_d_real = criterion(output_real, label_real)
            loss_d_fake = criterion(output_fake, label_fake)
            loss_d = loss_d_real + loss_d_fake

            # Backward and optimize discriminator
            discriminator.zero_grad()
            loss_d.backward()
            optimizer_d.step()

            # Train generator: max log(D(G(z)))
            noise = torch.randn(batch_size, input_dim_g).to(device)
            fake_data = generator(noise)

            # Forward pass fake data through discriminator
            output_fake = discriminator(fake_data)

            # Calculate the loss for generator
            loss_g = criterion(output_fake, label_real)

            # Backward and optimize generator
            generator.zero_grad()
            loss_g.backward()
            optimizer_g.step()

        print(f"Epoch [{epoch+1}/{num_epochs}], Loss D: {loss_d.item():.4f}, Loss G: {loss_g.item():.4f}")

    return generator


if __name__ == '__main__':
    # Set the seed value for reproducibility
    torch.manual_seed(42)

    # Load iris dataset and create the dataloader
    train_loader = load_data()

    # Initialize generator and discriminator
    generator = Generator(input_dim_g, output_dim_g).to(device)
    discriminator = Discriminator(input_dim_d).to(device)

    # Train the GAN model
    trained_generator = train(generator, discriminator, train_loader)

    # Generate fake data and convert it back to original scale
    with torch.no_grad():
        noise = torch.randn(100, input_dim_g).to(device)
        fake_data = trained_generator(noise)
        fake_data = fake_data.cpu().numpy()
        # Convert the normalized data back to the original scale
        stats = np.load("./data/iris_stats.npz")
        fake_data = fake_data * stats['std'] + stats['mean']

    # Save the generated data
    np.savetxt("./data/generated_data.csv", fake_data, delimiter=",")

在此示例代码中，我们使用 PyTorch 实现了简单的 GAN 模型。该模型由一个生成器和一个判别器组成。我们通过一个 4 维的鸢尾花数据集用于训练该 GAN 模型，并生成具有相同数据分布的合成数据集。

需要注意的是，在上述示例代码中，我们尚未对数据集进行任何增强。如果需要进行数据增强，则可以使用 PyTorch 的 `transforms` 模块进行数据增强操作。例如，以下代码演示了如何使用 `transforms` 模块进行数据增强操作：

transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(degrees=45),
    transforms.ToTensor(),
])

这里，我们使用了随机水平翻转 (`RandomHorizontalFlip`) 和随机旋转 (`RandomRotation`) 进行数据增强。最后，我们使用 `ToTensor` 将数据转换为张量形式，以便将其传递给 GAN 模型进行训练。