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

以下是使用PyTorch实现ACGAN自动生成动漫头像的代码：

首先，我们需要导入必要的库：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt

接下来，我们可以定义一些超参数：

# 超参数
batch_size = 128 # 批量大小
z_dim = 100 # 噪声维度
num_epochs = 200 # 训练轮数
lr = 0.0002 # 学习率
beta1 = 0.5 # Adam优化器的一项
num_classes = 10 # 类别数量

我们还需要定义生成器和判别器网络。这里我们采用类似DCGAN的架构，但是为了实现ACGAN，我们需要在判别器的输出上添加一个类别标签的预测。同时，我们也需要在生成器的输入上添加一个类别标签的条件。

# 定义生成器网络
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(z_dim + num_classes, 256, 7, 1, 0, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 3, 4, 2, 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(-1, z_dim + num_classes, 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, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, 4, 2, 1, bias=False),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 1, 7, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, img, labels):
        d_in = torch.cat((img, self.label_emb(labels)), 1)
        validity = self.model(d_in)
        return validity.view(-1)

然后，我们需要定义损失函数和优化器：

# 定义损失函数
adversarial_loss = nn.BCELoss()
auxiliary_loss = nn.CrossEntropyLoss()

# 定义生成器和判别器
generator = Generator()
discriminator = Discriminator()

# 定义优化器
optimizer_G = optim.Adam(generator.parameters(), lr=lr, betas=(beta1, 0.999))
optimizer_D = optim.Adam(discriminator.parameters(), lr=lr, betas=(beta1, 0.999))

接下来，我们需要加载数据集。这里我们使用了Anime Faces数据集，它包含了64x64大小的动漫头像图片。

# 加载数据集
transform = transforms.Compose([
    transforms.Resize(64),
    transforms.CenterCrop(64),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_dataset = datasets.ImageFolder('data/anime_faces', transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

最后，我们可以开始训练模型了：

# 训练模型
for epoch in range(num_epochs):
    for i, (imgs, labels) in enumerate(train_loader):

        # 确保batch大小是正确的
        if imgs.shape[0] != batch_size:
            continue

        # 生成噪声
        z = torch.randn((batch_size, z_dim))
        labels = labels.type(torch.LongTensor)

        # 训练判别器
        optimizer_D.zero_grad()

        real_imgs = imgs.cuda()
        real_labels = labels.cuda()
        fake_labels = torch.randint(0, num_classes, (batch_size,)).cuda()

        fake_imgs = generator(z.cuda(), fake_labels)

        real_validity = discriminator(real_imgs, real_labels)
        fake_validity = discriminator(fake_imgs.detach(), fake_labels)

        d_loss_real = adversarial_loss(real_validity, torch.ones_like(real_validity))
        d_loss_fake = adversarial_loss(fake_validity, torch.zeros_like(fake_validity))
        d_loss = (d_loss_real + d_loss_fake) / 2

        d_loss.backward()
        optimizer_D.step()

        # 训练生成器
        optimizer_G.zero_grad()

        gen_labels = torch.randint(0, num_classes, (batch_size,)).cuda()

        gen_imgs = generator(z.cuda(), gen_labels)
        validity = discriminator(gen_imgs, gen_labels)

        g_loss = adversarial_loss(validity, torch.ones_like(validity)) + auxiliary_loss(gen_labels, gen_labels)

        g_loss.backward()
        optimizer_G.step()

        # 打印损失
        if i % 10 == 0:
            print("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]" % (epoch, num_epochs, i, len(train_loader), d_loss.item(), g_loss.item()))

    # 每个epoch结束，生成一些示例图片
    with torch.no_grad():
        z = torch.randn((batch_size, z_dim)).cuda()
        labels = torch.arange(num_classes).repeat(batch_size // num_classes).cuda()
        gen_imgs = generator(z, labels).cpu().detach()

        # 将图片转换为numpy数组
        gen_imgs = gen_imgs.numpy().transpose((0, 2, 3, 1))

        # 反归一化
        gen_imgs = (gen_imgs + 1) / 2

        # 绘制示例图片
        fig, axs = plt.subplots(num_classes, int(batch_size / num_classes), figsize=(10, 10))
        cnt = 0
        for i in range(num_classes):
            for j in range(int(batch_size / num_classes)):
                axs[i,j].imshow(gen_imgs[cnt,:,:,:])
                axs[i,j].axis('off')
                cnt += 1
        fig.savefig("images/%d.png" % epoch)
        plt.close()

在训练过程中，我们每个epoch结束后都会生成一些示例图片，并保存到`images`目录下。下面是一些示例图片的样子：