用torch写一个GAN网络生成三维数据

以下一个简单的GAN网络用于生成三维数据的示例代码，使用PyTorch框架实现：

import torch
import torch.nn as nn
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

# 生成器网络
class Generator(nn.Module):
    def __init__(self, input_size, output_size):
        super(Generator, self).__init__()
        self.fc1 = nn.Linear(input_size, 128)
        self.fc2 = nn.Linear(128, 256)
        self.fc3 = nn.Linear(256, output_size)
        self.relu = nn.LeakyReLU(0.2)
        self.tanh = nn.Tanh()

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.fc3(out)
        out = self.tanh(out)
        return out

# 判别器网络
class Discriminator(nn.Module):
    def __init__(self, input_size):
        super(Discriminator, self).__init__()
        self.fc1 = nn.Linear(input_size, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 1)
        self.relu = nn.LeakyReLU(0.2)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.fc3(out)
        out = self.sigmoid(out)
        return out

# 定义参数
input_size = 100
output_size = 3
num_epochs = 20000
batch_size = 64
learning_rate = 0.0002

# 实例化网络和优化器
G = Generator(input_size, output_size)
D = Discriminator(output_size)
criterion = nn.BCELoss()
G_optimizer = torch.optim.Adam(G.parameters(), lr=learning_rate)
D_optimizer = torch.optim.Adam(D.parameters(), lr=learning_rate)

# 生成随机噪声函数
def noise(size):
    n = torch.randn(size, input_size)
    return n

# 定义真实数据
def real_data_generator(size):
    x = np.random.uniform(-1, 1, size)
    y = np.random.uniform(-1, 1, size)
    z = np.random.uniform(-1, 1, size)
    data = np.stack((x, y, z), axis=1)
    return torch.from_numpy(data).float()

# 训练GAN网络
for epoch in range(num_epochs):
    # 判别器训练
    for _ in range(5):
        # 生成假数据
        z = noise(batch_size)
        fake_data = G(z)

        # 计算损失
        D_real = D(real_data_generator(batch_size))
        D_fake = D(fake_data)
        D_loss = criterion(D_real, torch.ones(batch_size)) + criterion(D_fake, torch.zeros(batch_size))

        # 反向传播和优化
        D_optimizer.zero_grad()
        D_loss.backward()
        D_optimizer.step()

    # 生成器训练
    z = noise(batch_size)
    fake_data = G(z)
    D_fake = D(fake_data)
    G_loss = criterion(D_fake, torch.ones(batch_size))

    # 反向传播和优化
    G_optimizer.zero_grad()
    G_loss.backward()
    G_optimizer.step()

    # 打印损失
    if epoch % 1000 == 0:
        print(f"Epoch {epoch}, Generator loss: {G_loss.item()}, Discriminator loss: {D_loss.item()}")

# 生成数据并可视化
z = noise(1000)
generated_data = G(z).detach().numpy()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(generated_data[:, 0], generated_data[:, 1], generated_data[:, 2], s=5)
plt.show()

在训练过程中，我们使用随机噪声作为生成器的输入，生成器负责将噪声转换为三维数据。判别器负责判断输入的数据是真实的还是生成的。GAN网络的训练是一个对抗过程，生成器和判别器相互竞争，直到生成器可以生成逼真的三维数据。