GAN-生成对抗网络-生成人脸图像-CNN

Posted 2022-03-22 挂科难

数据集

香港中文大学，CeleA数据集，http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
总共202599张面部图片。

制作HDF5压缩格式文件

分层数据格式（hierarchical data format）是一种成熟的，开源的压缩数据格式，专门用于存储非常大量的数据。
在《PyTorch 生成对抗网络编程》[英]塔克里·拉希德一书中将数据处理成HDF5格式运行代码，本来想重构一下的。然后发现如果用ImgaeFloder直接读取的话基本要重构所有代码。还是制作HDF5吧！

import os
import zipfile

import h5py
import imageio

hdf5_file = './celeba_aligned_small.h5py'

total_images = 202599

with h5py.File(hdf5_file, 'w') as hf:  # 打开h5py文件，文件不存在则会创建文件

    count = 0

    with zipfile.ZipFile('img_align_celeba.zip', 'r') as zf:
    # 这个压缩文件里是一个文件夹img_align_celeba文件夹中有200000多张图片
        for i in zf.namelist():  # zf.namelist()返还压缩文件中的文件列表名
        # zf.namelist()[0]是'img_align_celeba/' 即文件夹路径
        # zf.namelist()[1]是'img_align_celeba/000001.jpg' 即文件夹下的文件路径
            if i[-4:] == '.jpg':
                ofile = zf.extract(i)  # 解压单个文件至ofile中
                # 默认解压在当前文件夹即在'./'路径下创建img_align_celeba文件夹，把图片(i)放入文件夹中
                # ofile是解压后图片(i)的相对地址是一个字符串
                img = imageio.imread(ofile)
                # 使用imageio.imread读取图片，此时img打印出来是一个数组
                os.remove(ofile)  # 用完即弃
                # 删除图片，不占存储空间

                hf.create_dataset('img_align_celeba/'+str(count)+'.jpg',
                                  data=img, compression='gzip', compression_opts=9)
                # compression是压缩方式, compression_opts是压缩程度的参数
                # 在celeba_aligned_small.h5py文件中生成组img_align_celeba，在组中保存img数组

                count = count + 1
                if count % 1000 == 0:
                    print('images done ...', count)

                if count == total_images:  # 取前所有图片
                    break

生成人脸

上篇博文使用的是MLP（全连接神经网络），这次生成人脸换成了CNN（卷积神经网络）。并采用GPU加速

import h5py
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset
import torch
import torch.nn as nn


def crop_centre(img, new_width, new_height):
    height, width, _ = img.shape
    startx = width // 2 - new_width // 2
    starty = height // 2 - new_height // 2
    return img[starty:starty + new_height, startx:startx + new_width, :]


def generate_random_image(size):
    random_data = torch.rand(size)
    return torch.cuda.FloatTensor(random_data)


def generate_random_seed(size):
    random_data = torch.randn(size)
    return random_data


class CelebADataset(Dataset):
    def __init__(self, file):
        self.file_object = h5py.File(file, 'r')
        self.dataset = self.file_object['img_align_celeba']

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        if index >= len(self.dataset):
            raise IndexError
        img = np.array(self.dataset[str(index) + '.jpg'])
        img = crop_centre(img, 128, 128)
        return torch.cuda.FloatTensor(img).permute(2, 0, 1).view(1, 3, 128, 128) / 255.0

    def plot_image(self, index):
        img = np.array(self.dataset[str(index) + '.jpg'])
        img = crop_centre(img, 128, 128)
        plt.imshow(img, interpolation='nearest')
        plt.show()


# 构建鉴别器
class Discriminator(nn.Module):
    def __init__(self):
        # 初始化父类
        super(Discriminator, self).__init__()

        self.feature = nn.Sequential(
            nn.Conv2d(3, 256, kernel_size=8, stride=2),
            nn.BatchNorm2d(256),
            nn.GELU(),

            nn.Conv2d(256, 256, kernel_size=8, stride=2),
            nn.BatchNorm2d(256),
            nn.GELU(),

            nn.Conv2d(256, 3, kernel_size=8, stride=2),
            nn.GELU(),
        )
        self.classifier = nn.Sequential(
            nn.Linear(3 * 10 * 10, 1),
            nn.Sigmoid()
        )

        self.loss_function = nn.BCELoss()

        # 创建优化器
        self.optimiser = torch.optim.Adam(self.parameters(), lr=0.01)

        self.counter = 0
        self.progress = []

    def forward(self, inputs):
        x = self.feature(inputs)
        x = x.view(-1)
        x = self.classifier(x)
        return x

    def train(self, inputs, targets):
        outputs = self.forward(inputs)
        loss = self.loss_function(outputs, targets)

        # 每训练10此增加计数器
        self.counter += 1
        if self.counter % 10 == 0:
            self.progress.append(loss.item())

        if self.counter % 10000 == 0:
            print("counter = ", self.counter)

        # 清楚梯度，反向传播， 更新权重
        self.optimiser.zero_grad()
        loss.backward()
        self.optimiser.step()


# 构建生成器
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()

        self.linear = nn.Sequential(
            nn.Linear(100, 3*11*11),
            nn.GELU(),
        )

        self.feature = nn.Sequential(
            nn.ConvTranspose2d(3, 256, kernel_size=8, stride=2),
            nn.BatchNorm2d(256),
            nn.GELU(),

            nn.ConvTranspose2d(256, 256, kernel_size=8, stride=2),
            nn.BatchNorm2d(256),
            nn.GELU(),

            nn.ConvTranspose2d(256, 3, kernel_size=8, stride=2, padding=1),
            nn.BatchNorm2d(3),

            nn.Sigmoid()
        )

        # 创建优化器
        self.optimiser = torch.optim.Adam(self.parameters(), lr=0.01)

        self.counter = 0
        self.progress = []

    def forward(self, x):
        x = self.linear(x)
        x = x.view(1, 3, 11, 11)
        x = self.feature(x)
        return x

    def train(self, D, inputs, targets):  # 用鉴别器的损失来训练生成，discriminator

        g_output = self.forward(inputs)   # 生成器generator的输出(g_output)

        d_output = D.forward(g_output)  # 分类器discriminator的输出,输入生成器generator的输出

        loss = D.loss_function(d_output, targets)

        self.counter += 1
        if self.counter % 10 == 0:
            self.progress.append(loss.item())

        self.optimiser.zero_grad()
        loss.backward()
        self.optimiser.step()


if torch.cuda.is_available():
    torch.set_default_tensor_type(torch.cuda.FloatTensor)
    print('using cuda:', torch.cuda.get_device_name(0))

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 现在浮点类型默认cuda张量
# x = generate_random_image(2)
# print(x.device)

celeba_dataset = CelebADataset('celeba_aligned_small.h5py')
# celeba_dataset.plot_image(66) # 输出一张图试试
D = Discriminator()
G = Generator()
G.to(device)
D.to(device)
epoches = 1
for epoch in range(epoches):
    print('开始第', epoch+1, '轮', '*************'*4)
    for image_data in celeba_dataset:
        # 真实样本1，训练鉴别器
        D.train(image_data, torch.cuda.FloatTensor([1.0]))
        # 随机生成噪声，告诉鉴别器这是0
        D.train(generate_random_image((1, 3, 128, 128)), torch.cuda.FloatTensor([0.0]))

        # 训练生成器
        G.train(D, generate_random_seed(100), torch.cuda.FloatTensor([1.0]))

# torch.save(G, 'face-cnn-1.pth')

多运行几轮能得到不错的效果：

for i in range(6):
    output = G.forward(generate_random_seed(100))
    img = output.detach().permute(0,2,3,1).view(128,128,3).cpu().numpy()
    plt.subplot(2, 3, i+1)
    plt.imshow(img)
plt.show()

运行一轮我这1060跑了4个小时，设备好的炼丹师们可以多运行几轮试试。