昇思25天学习打卡营第24天 | DCGAN生成漫画头像

import mindspore as ms
from mindspore import nn, ops, context
from mindspore.common.initializer import Normal
import mindspore.dataset as ds
import mindspore.dataset.vision as vision
import matplotlib.pyplot as plt
import numpy as np
from download import download

# 设置环境
context.set_context(mode=context.GRAPH_MODE, device_target="CPU")

# 下载并解压数据集
url = "https://download.mindspore.cn/dataset/Faces/faces.zip"
path = download(url, "./faces", kind="zip", replace=True)

# 数据处理
batch_size = 128
image_size = 64
nc = 3
nz = 100
ngf = 64
ndf = 64
num_epochs = 3
lr = 0.0002
beta1 = 0.5

def create_dataset_imagenet(dataset_path):
    dataset = ds.ImageFolderDataset(dataset_path, num_parallel_workers=4, shuffle=True, decode=True)
    transforms = [
        vision.Resize(image_size),
        vision.CenterCrop(image_size),
        vision.HWC2CHW(),
        lambda x: ((x / 255).astype("float32"))
    ]
    dataset = dataset.project('image')
    dataset = dataset.map(transforms, 'image')
    dataset = dataset.batch(batch_size)
    return dataset

dataset = create_dataset_imagenet('./faces')

def plot_data(data):
    plt.figure(figsize=(10, 3), dpi=140)
    for i, image in enumerate(data[0][:30], 1):
        plt.subplot(3, 10, i)
        plt.axis("off")
        plt.imshow(image.transpose(1, 2, 0))
    plt.show()

sample_data = next(dataset.create_tuple_iterator(output_numpy=True))
plot_data(sample_data)

# 生成器
class Generator(nn.Cell):
    def __init__(self):
        super(Generator, self).__init__()
        self.generator = nn.SequentialCell(
            nn.Conv2dTranspose(nz, ngf * 8, 4, 1, 'valid', weight_init=Normal(0.02)),
            nn.BatchNorm2d(ngf * 8, gamma_init=Normal(1.0, 0.02)),
            nn.ReLU(),
            nn.Conv2dTranspose(ngf * 8, ngf * 4, 4, 2, 'pad', 1, weight_init=Normal(0.02)),
            nn.BatchNorm2d(ngf * 4, gamma_init=Normal(1.0, 0.02)),
            nn.ReLU(),
            nn.Conv2dTranspose(ngf * 4, ngf * 2, 4, 2, 'pad', 1, weight_init=Normal(0.02)),
            nn.BatchNorm2d(ngf * 2, gamma_init=Normal(1.0, 0.02)),
            nn.ReLU(),
            nn.Conv2dTranspose(ngf * 2, ngf, 4, 2, 'pad', 1, weight_init=Normal(0.02)),
            nn.BatchNorm2d(ngf, gamma_init=Normal(1.0, 0.02)),
            nn.ReLU(),
            nn.Conv2dTranspose(ngf, nc, 4, 2, 'pad', 1, weight_init=Normal(0.02)),
            nn.Tanh()
        )

    def construct(self, x):
        return self.generator(x)

generator = Generator()

# 判别器
class Discriminator(nn.Cell):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.discriminator = nn.SequentialCell(
            nn.Conv2d(nc, ndf, 4, 2, 'pad', 1, weight_init=Normal(0.02)),
            nn.LeakyReLU(0.2),
            nn.Conv2d(ndf, ndf * 2, 4, 2, 'pad', 1, weight_init=Normal(0.02)),
            nn.BatchNorm2d(ndf * 2, gamma_init=Normal(1.0, 0.02)),
            nn.LeakyReLU(0.2),
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 'pad', 1, weight_init=Normal(0.02)),
            nn.BatchNorm2d(ndf * 4, gamma_init=Normal(1.0, 0.02)),
            nn.LeakyReLU(0.2),
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 'pad', 1, weight_init=Normal(0.02)),
            nn.BatchNorm2d(ndf * 8, gamma_init=Normal(1.0, 0.02)),
            nn.LeakyReLU(0.2),
            nn.Conv2d(ndf * 8, 1, 4, 1, 'valid', weight_init=Normal(0.02)),
        )
        self.adv_layer = nn.Sigmoid()

    def construct(self, x):
        out = self.discriminator(x)
        out = out.reshape(out.shape[0], -1)
        return self.adv_layer(out)

discriminator = Discriminator()

# 损失函数和优化器
adversarial_loss = nn.BCELoss(reduction='mean')
optimizer_D = nn.Adam(discriminator.trainable_params(), learning_rate=lr, beta1=beta1)
optimizer_G = nn.Adam(generator.trainable_params(), learning_rate=lr, beta1=beta1)
optimizer_G.update_parameters_name('optim_g.')
optimizer_D.update_parameters_name('optim_d.')

def generator_forward(real_imgs, valid):
    z = ops.standard_normal((real_imgs.shape[0], nz, 1, 1))
    gen_imgs = generator(z)
    g_loss = adversarial_loss(discriminator(gen_imgs), valid)
    return g_loss, gen_imgs

def discriminator_forward(real_imgs, gen_imgs, valid, fake):
    real_loss = adversarial_loss(discriminator(real_imgs), valid)
    fake_loss = adversarial_loss(discriminator(gen_imgs), fake)
    d_loss = (real_loss + fake_loss) / 2
    return d_loss

grad_generator_fn = ms.value_and_grad(generator_forward, None, optimizer_G.parameters, has_aux=True)
grad_discriminator_fn = ms.value_and_grad(discriminator_forward, None, optimizer_D.parameters)

@ms.jit
def train_step(imgs):
    valid = ops.ones((imgs.shape[0], 1), mindspore.float32)
    fake = ops.zeros((imgs.shape[0], 1), mindspore.float32)

    (g_loss, gen_imgs), g_grads = grad_generator_fn(imgs, valid)
    optimizer_G(g_grads)
    d_loss, d_grads = grad_discriminator_fn(imgs, gen_imgs, valid, fake)
    optimizer_D(d_grads)

    return g_loss, d_loss, gen_imgs

G_losses = []
D_losses = []
image_list = []

total = dataset.get_dataset_size()
for epoch in range(num_epochs):
    generator.set_train()
    discriminator.set_train()
    for i, (imgs, ) in enumerate(dataset.create_tuple_iterator()):
        g_loss, d_loss, gen_imgs = train_step(imgs)
        if i % 100 == 0 or i == total - 1:
            print('[%2d/%d][%3d/%d]   Loss_D:%7.4f  Loss_G:%7.4f' % (
                epoch + 1, num_epochs, i + 1, total, d_loss.asnumpy(), g_loss.asnumpy()))
        D_losses.append(d_loss.asnumpy())
        G_losses.append(g_loss.asnumpy())

    generator.set_train(False)
    fixed_noise = ops.standard_normal((batch_size, nz, 1, 1))
    img = generator(fixed_noise)
    image_list.append(img.transpose(0, 2, 3, 1).asnumpy())

    mindspore.save_checkpoint(generator, "./generator.ckpt")
    mindspore.save_checkpoint(discriminator, "./discriminator.ckpt")

# 可视化损失
plt.figure(figsize=(10,5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(G_losses, label="G")
plt.plot(D_losses, label="D")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.show()

# 可视化生成的图像
for i, imgs in enumerate(image_list):
    plt.figure(figsize=(10, 10))
    for j in range(25):
        plt.subplot(5, 5, j+1)
        plt.axis('off')
        plt.imshow((imgs[j] + 1) / 2)
    plt.show()

确保你已经安装了所需的包，包括 MindSpore、Matplotlib 和 Numpy。你可以根据需要调整超参数和训练细节。运行这段代码后，你将能够生成动漫头像并可视化训练过程中的损失变化。