【生成式对抗网络】GANs在数据生成、艺术创作，以及在增强现实和虚拟现实中的应用

import torch
from torch import nn
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader
from torchvision.utils import save_image
# 定义生成器和判别器
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            nn.ConvTranspose2d(100, 128, 7),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.ConvTranspose2d(128, 64, 4, 2, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 1, 4, 2, 1),
            nn.Tanh()
        )
    def forward(self, input):
        return self.main(input)
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            nn.Conv2d(1, 64, 4, 2, 1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, 4, 2, 1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 1, 7),
            nn.Sigmoid()
        )
    def forward(self, input):
        return self.main(input).view(-1)
# 设置超参数
batch_size = 128
lr = 0.0002
epochs = 20
# 加载MNIST数据集
train_data = MNIST('.', download=True, transform=ToTensor())
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
# 初始化模型
device = 'cuda' if torch.cuda.is_available() else 'cpu'
generator = Generator().to(device)
discriminator = Discriminator().to(device)
criterion = nn.BCELoss()
# 训练循环
for epoch in range(epochs):
    for i, (real_images, _) in enumerate(train_loader):
        # 训练判别器
        real_images = real_images.to(device)
        real_labels = torch.ones(real_images.size(0)).to(device)
        fake_labels = torch.zeros(real_images.size(0)).to(device)
        
        # 生成假图像
        noise = torch.randn(real_images.size(0), 100, 1, 1).to(device)
        fake_images = generator(noise)
        
        # 判别器损失
        real_loss = criterion(discriminator(real_images), real_labels)
        fake_loss = criterion(discriminator(fake_images.detach()), fake_labels)
        d_loss = real_loss + fake_loss
        
        # 更新判别器参数
        discriminator.zero_grad()
        d_loss.backward()
        optimizer_d.step()
        
        # 训练生成器
        noise = torch.randn(real_images.size(0), 100, 1, 1).to(device)
        fake_images = generator(noise)
        
        # 生成器损失
        g_loss = criterion(discriminator(fake_images), real_labels)
        
        # 更新生成器参数
        generator.zero_grad()
        g_loss.backward()
        optimizer_g.step()
        
        # 打印损失
        print(f"Epoch [{epoch}/{epochs}], Step [{i}/{len(train_loader)}], d_loss: {d_loss.item()}, g_loss: {g_loss.item()}")
        
        # 每隔一定步数保存生成的图像
        if i % 200 == 0:
            save_image(fake_images.data[:25], f'sample_{epoch}_{i}.png', nrow=5, normalize=True)

from cycle_gan import CycleGANModel
from data import get_dataloader
# 加载数据
dataloader = get_dataloader('horses', 'zebras')
# 初始化 CycleGAN 模型
model = CycleGANModel()
# 训练模型
model.train(dataloader, epochs=100)
# 测试模型
model.test(dataloader)