ACGAN自动生成动漫头像PyTorch 代码

以下是基于 PyTorch 的 ACGAN 自动生成动漫头像的代码：

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 torchvision.datasets import ImageFolder
from tqdm import tqdm

# 定义超参数
num_epochs = 200
batch_size = 64
learning_rate = 0.0002
latent_size = 100
num_classes = 10
image_size = 64

# 加载数据集
dataset = ImageFolder(root='./data', transform=transforms.Compose([
    transforms.Resize(image_size),
    transforms.CenterCrop(image_size),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]))
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# 定义生成器模型
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.label_emb = nn.Embedding(num_classes, num_classes)
        self.model = nn.Sequential(
            nn.ConvTranspose2d(latent_size + num_classes, 512, kernel_size=4, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1, bias=False),
            nn.Tanh()
        )

    def forward(self, noise, labels):
        gen_input = torch.cat((self.label_emb(labels), noise), -1)
        gen_input = gen_input.view(gen_input.size(0), gen_input.size(1), 1, 1)
        img = self.model(gen_input)
        return img

# 定义判别器模型
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.label_emb = nn.Embedding(num_classes, num_classes)
        self.model = nn.Sequential(
            nn.Conv2d(3 + num_classes, 64, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, img, labels):
        dis_input = torch.cat((img, self.label_emb(labels)), -1)
        output = self.model(dis_input)
        return output.view(-1, 1).squeeze(1)

# 初始化模型和优化器
generator = Generator()
discriminator = Discriminator()
generator.cuda()
discriminator.cuda()
adversarial_loss = nn.BCELoss()
categorical_loss = nn.CrossEntropyLoss()
optimizer_G = optim.Adam(generator.parameters(), lr=learning_rate, betas=(0.5, 0.999))
optimizer_D = optim.Adam(discriminator.parameters(), lr=learning_rate, betas=(0.5, 0.999))

# 训练模型
for epoch in range(num_epochs):
    for i, (images, labels) in tqdm(enumerate(dataloader)):
        # 训练判别器
        discriminator.zero_grad()
        real_images = images.cuda()
        real_labels = labels.cuda()
        batch_size = real_images.size(0)
        real_outputs = discriminator(real_images, real_labels)
        real_loss = adversarial_loss(real_outputs, torch.ones(batch_size).cuda())
        real_loss.backward()

        noise = torch.randn(batch_size, latent_size).cuda()
        fake_labels = torch.randint(0, num_classes, (batch_size,)).cuda()
        fake_images = generator(noise, fake_labels)
        fake_outputs = discriminator(fake_images.detach(), fake_labels)
        fake_loss = adversarial_loss(fake_outputs, torch.zeros(batch_size).cuda())
        fake_loss.backward()
        d_loss = real_loss + fake_loss
        optimizer_D.step()

        # 训练生成器
        generator.zero_grad()
        noise = torch.randn(batch_size, latent_size).cuda()
        fake_labels = torch.randint(0, num_classes, (batch_size,)).cuda()
        fake_images = generator(noise, fake_labels)
        outputs = discriminator(fake_images, fake_labels)
        g_loss = adversarial_loss(outputs, torch.ones(batch_size).cuda())
        g_loss.backward()
        optimizer_G.step()

        if (i+1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], D_loss: {:.4f}, G_loss: {:.4f}'
                  .format(epoch+1, num_epochs, i+1, len(dataloader), d_loss.item(), g_loss.item()))

    # 保存生成的图片
    with torch.no_grad():
        noise = torch.randn(10, latent_size).cuda()
        labels = torch.arange(0, 10).unsqueeze(1).cuda()
        fake_images = generator(noise, labels).detach().cpu()
        save_image(fake_images, './generated_images_{}.png'.format(epoch+1), nrow=10, normalize=True)

# 保存模型
torch.save(generator.state_dict(), 'generator.pth')
torch.save(discriminator.state_dict(), 'discriminator.pth')

在这个代码中，我们使用了 PyTorch 的 nn.Module 和 nn.Sequential 构建了生成器和判别器模型。同时，我们使用了 DataLoader 加载了动漫头像数据集，并定义了超参数和优化器。最后，我们使用 for 循环进行模型的训练，并保存了生成的图片和模型。