import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from sklearn.datasets import load_iris
import warnings
# 忽略不必要的警告信息
warnings.filterwarnings("ignore")
# --------------------------
# 1. 配置训练参数与设备
# --------------------------
# 潜在空间维度（生成器的输入维度）
latent_dim = 10
# 训练总轮数（GAN通常需要较多迭代才能收敛）
train_epochs = 10000
# 批次大小（根据数据集规模调整）
batch_size = 32
# 学习率（控制参数更新幅度）
learning_rate = 0.0002
# Adam优化器的动量参数（影响收敛稳定性）
beta1 = 0.5
# 自动选择运算设备（优先GPU，没有则用CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"当前使用设备: {device}")
# --------------------------
# 2. 数据加载与预处理
# --------------------------
# 加载鸢尾花数据集
iris_dataset = load_iris()
# 提取特征数据和标签
features = iris_dataset.data
labels = iris_dataset.target
# 只选取Setosa类别（标签为0）的数据进行训练
setosa_features = features[labels == 0]
# 将数据缩放到[-1, 1]区间（配合生成器的Tanh输出激活）
scaler = MinMaxScaler(feature_range=(-1, 1))
scaled_features = scaler.fit_transform(setosa_features)
# 转换为PyTorch张量并创建数据加载器
# 注意：必须转为float类型才能与模型参数兼容
data_tensor = torch.from_numpy(scaled_features).float()
dataset = TensorDataset(data_tensor)
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
# 打印数据基本信息
print(f"训练样本数量: {len(scaled_features)}")
print(f"特征维度: {scaled_features.shape[1]}") # 鸢尾花数据集固定为4维特征
# --------------------------
# 3. 定义生成器和判别器
# --------------------------
class Generator(nn.Module):
"""生成器：将随机噪声转换为模拟的鸢尾花特征数据"""
def __init__(self):
super(Generator, self).__init__()
# 简单的全连接网络结构
self.net = nn.Sequential(
nn.Linear(latent_dim, 16), # 从潜在空间映射到16维
nn.ReLU(), # 激活函数增加非线性
nn.Linear(16, 32), # 进一步映射到32维
nn.ReLU(),
nn.Linear(32, 4), # 输出4维特征（与真实数据一致）
nn.Tanh() # 确保输出在[-1, 1]范围内
)
def forward(self, x):
# 前向传播：输入噪声，输出生成的数据
return self.net(x)
class Discriminator(nn.Module):
"""判别器：区分输入数据是真实样本还是生成器伪造的"""
def __init__(self):
super(Discriminator, self).__init__()
# 简单的全连接网络结构
self.net = nn.Sequential(
nn.Linear(4, 32), # 输入4维特征
nn.LeakyReLU(0.2), # LeakyReLU避免梯度消失问题
nn.Linear(32, 16), # 压缩到16维
nn.LeakyReLU(0.2),
nn.Linear(16, 1), # 输出单个概率值
nn.Sigmoid() # 将输出压缩到[0,1]（表示真实数据的概率）
)
def forward(self, x):
# 前向传播：输入数据，输出判断概率
return self.net(x)
# 初始化模型并移动到运算设备
generator = Generator().to(device)
discriminator = Discriminator().to(device)
# 打印模型结构
print("\n生成器结构:")
print(generator)
print("\n判别器结构:")
print(discriminator)
# --------------------------
# 4. 配置训练组件
# --------------------------
# 定义损失函数（二元交叉熵，适合二分类问题）
criterion = nn.BCELoss()
# 定义优化器（分别优化生成器和判别器）
gen_optimizer = optim.Adam(generator.parameters(), lr=learning_rate, betas=(beta1, 0.999))
dis_optimizer = optim.Adam(discriminator.parameters(), lr=learning_rate, betas=(beta1, 0.999))
# --------------------------
# 5. 开始训练
# --------------------------
print("\n--- 训练开始 ---")
for epoch in range(train_epochs):
# 遍历数据加载器中的每一批次
for batch_idx, (real_data,) in enumerate(data_loader):
# 将真实数据移动到运算设备
real_data = real_data.to(device)
current_batch_size = real_data.size(0) # 获取当前批次的实际样本数（最后一批可能不满）
# 创建标签：真实数据标为1，生成数据标为0
real_labels = torch.ones(current_batch_size, 1).to(device)
fake_labels = torch.zeros(current_batch_size, 1).to(device)
# --------------------
# 训练判别器
# --------------------
dis_optimizer.zero_grad() # 清空判别器的梯度缓存
# 1. 用真实数据训练
real_output = discriminator(real_data)
# 计算真实数据的损失（希望判别器能认出真实数据）
loss_real = criterion(real_output, real_labels)
# 2. 用生成的数据训练
# 生成随机噪声（作为生成器的输入）
noise = torch.randn(current_batch_size, latent_dim).to(device)
# 生成假数据，并阻断梯度流向生成器（避免影响生成器参数）
fake_data = generator(noise).detach()
fake_output = discriminator(fake_data)
# 计算假数据的损失（希望判别器能认出假数据）
loss_fake = criterion(fake_output, fake_labels)
# 总损失反向传播并更新判别器参数
dis_loss = loss_real + loss_fake
dis_loss.backward()
dis_optimizer.step()
# --------------------
# 训练生成器
# --------------------
gen_optimizer.zero_grad() # 清空生成器的梯度缓存
# 重新生成假数据（这次需要计算生成器的梯度）
noise = torch.randn(current_batch_size, latent_dim).to(device)
fake_data = generator(noise)
fake_output = discriminator(fake_data)
# 生成器的损失：希望判别器把假数据当成真的（所以标签用real_labels）
gen_loss = criterion(fake_output, real_labels)
gen_loss.backward()
gen_optimizer.step()
# 每1000轮打印一次训练状态
if (epoch + 1) % 1000 == 0:
print(
f"轮次 [{epoch+1}/{train_epochs}], "
f"判别器损失: {dis_loss.item():.4f}, "
f"生成器损失: {gen_loss.item():.4f}"
)
print("\n--- 训练完成 ---")