wgan训练自己的数据集

以下是一个简单的 WGAN（Wasserstein GAN）的代码示例：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.autograd import Variable

# 定义生成器模型
class Generator(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Generator, self).__init__()
        self.fc = nn.Linear(input_dim, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, output_dim)
        self.tanh = nn.Tanh()

    def forward(self, x):
        x = self.fc(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.tanh(x)
        return x

# 定义判别器模型
class Discriminator(nn.Module):
    def __init__(self, input_dim):
        super(Discriminator, self).__init__()
        self.fc = nn.Linear(input_dim, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 1)
        
    def forward(self, x):
        x = self.fc(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 定义训练函数
def train_wgan(generator, discriminator, dataloader, num_epochs):
    optimizer_G = optim.RMSprop(generator.parameters(), lr=0.00005)
    optimizer_D = optim.RMSprop(discriminator.parameters(), lr=0.00005)
    criterion = nn.MSELoss()

    for epoch in range(num_epochs):
        for i, data in enumerate(dataloader):
            real_data = data[0]
            batch_size = real_data.size(0)
            real_data = Variable(real_data.view(batch_size, -1))
            
            # 训练判别器
            optimizer_D.zero_grad()
            z = Variable(torch.randn(batch_size, 100))
            fake_data = generator(z)
            real_pred = discriminator(real_data)
            fake_pred = discriminator(fake_data.detach())
            d_loss = -torch.mean(real_pred) + torch.mean(fake_pred)
            d_loss.backward()
            optimizer_D.step()
            
            # 截断权重
            for p in discriminator.parameters():
                p.data.clamp_(-0.01, 0.01)
            
            # 训练生成器
            if i % 5 == 0:
                optimizer_G.zero_grad()
                z = Variable(torch.randn(batch_size, 100))
                fake_data = generator(z)
                fake_pred = discriminator(fake_data)
                g_loss = -torch.mean(fake_pred)
                g_loss.backward()
                optimizer_G.step()

        print(f"Epoch [{epoch+1}/{num_epochs}], d_loss: {d_loss.item():.4f}, g_loss: {g_loss.item():.4f}")

# 定义数据集和数据加载器
transform = trans