输入为11*11大小数据集用wgan生成11*11灰度图代码

下面是一个简单的WGAN-GP模型的代码示例，可以用于生成11x11大小的灰度图像。这里使用PyTorch实现。

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torch.utils.data import DataLoader
from torch.autograd import grad

# 定义生成器和判别器的网络结构
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.fc = nn.Linear(100, 256)
        self.conv = nn.Sequential(
            nn.ConvTranspose2d(16, 8, 4, stride=2, padding=1),
            nn.BatchNorm2d(8),
            nn.ReLU(),
            nn.ConvTranspose2d(8, 1, 4, stride=2, padding=1),
            nn.Tanh()
        )

    def forward(self, z):
        x = self.fc(z)
        x = x.view(-1, 16, 4, 4)
        x = self.conv(x)
        return x

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 8, 4, stride=2, padding=1),
            nn.BatchNorm2d(8),
            nn.LeakyReLU(),
            nn.Conv2d(8, 16, 4, stride=2, padding=1),
            nn.BatchNorm2d(16),
            nn.LeakyReLU()
        )
        self.fc = nn.Linear(256, 1)

    def forward(self, x):
        x = self.conv(x)
        x = x.view(-1, 256)
        x = self.fc(x)
        return x

# 定义WGAN-GP的损失函数
def wgan_gp_loss(real, fake, discriminator, device):
    # 计算判别器对真实图像和生成图像的输出
    real_out = discriminator(real)
    fake_out = discriminator(fake)

    # 计算WGAN-GP损失
    d_loss = fake_out.mean() - real_out.mean()
    epsilon = torch.rand(real.shape[0], 1, 1, 1).to(device)
    interpolated = epsilon * real + (1 - epsilon) * fake
    interpolated_out = discriminator(interpolated)
    gradients = grad(outputs=interpolated_out, inputs=interpolated,
                     grad_outputs=torch.ones_like(interpolated_out),
                     create_graph=True, retain_graph=True, only_inputs=True)[0]
    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * 10
    d_loss += gradient_penalty

    return d_loss

# 设置训练参数和超参数
batch_size = 64
lr = 0.0001
z_dim = 100
n_epochs = 200
clip_value = 0.01
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载数据集
transform = transforms.Compose([
    transforms.Resize(11),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])
dataset = DataLoader(torchvision.datasets.MNIST(root='./data', train=True, transform=transform, download=True),
                     batch_size=batch_size, shuffle=True)

# 初始化生成器和判别器
generator = Generator().to(device)
discriminator = Discriminator().to(device)

# 定义优化器
g_optimizer = optim.Adam(generator.parameters(), lr=lr, betas=(0.5, 0.999))
d_optimizer = optim.Adam(discriminator.parameters(), lr=lr, betas=(0.5, 0.999))

# 训练WGAN-GP模型
for epoch in range(n_epochs):
    for i, (real_images, _) in enumerate(dataset):
        real_images = real_images.to(device)

        # 训练判别器
        for j in range(5):
            # 生成随机噪声
            z = torch.randn(real_images.shape[0], z_dim).to(device)

            # 生成假图像
            fake_images = generator(z)

            # 更新判别器
            d_optimizer.zero_grad()
            d_loss = wgan_gp_loss(real_images, fake_images, discriminator, device)
            d_loss.backward()
            d_optimizer.step()

            # 截断判别器的参数
            for p in discriminator.parameters():
                p.data.clamp_(-clip_value, clip_value)

        # 训练生成器
        z = torch.randn(real_images.shape[0], z_dim).to(device)
        fake_images = generator(z)
        g_optimizer.zero_grad()
        g_loss = -discriminator(fake_images).mean()
        g_loss.backward()
        g_optimizer.step()

        # 输出训练信息
        if i % 100 == 0:
            print(f"Epoch [{epoch}/{n_epochs}] Batch [{i}/{len(dataset)}] D loss: {d_loss:.4f} | G loss: {g_loss:.4f}")

    # 保存生成的图像
    with torch.no_grad():
        z = torch.randn(1, z_dim).to(device)
        fake_image = generator(z).squeeze()
        save_image(fake_image, f"images/{epoch}.png", normalize=True)

这个代码示例中，我们定义了一个生成器和一个判别器，它们分别用于生成和判别11x11的灰度图像。在训练过程中，我们使用了WGAN-GP损失函数，并使用Adam优化器对生成器和判别器进行优化。在每个epoch结束时，我们生成一个随机噪声向量，并使用生成器生成一个假图像，并将其保存为一个PNG文件。