文章目录

  • 1、导入需要的库
  • 2、导入数据集
  • 2.1 读入文件夹中的图片
  • 2.2 初始化导入的图片
  • 2.3 构建数据集对象
  • 3、构建DCGAN网络
  • 3.1 建立生成器
  • 3.2 建立判别器
  • 3.3 实例化判别器和生成器
  • 4、建立损失函数
  • 4.1 生成器损失函数
  • 4.2 判别器损失函数
  • 5、初始化优化器
  • 6、定义梯度下降过程
  • 7、将生成的多张图像放到一个图里并保存
  • 8、训练


1、导入需要的库

import os
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import glob
from PIL import Image
image_size = 64
SHUFFLE_SIZE = 1000
batch_size = 64
z_dim = 100 # 隐藏向量z的长度
img_shape = (image_size, image_size, 3)

2、导入数据集

2.1 读入文件夹中的图片

PATH = '.\\faces/'
X_train = tf.data.Dataset.list_files(PATH+'*.jpg')
img_path = glob.glob(r'.\faces\*.jpg')
print('images num:', len(img_path))
images num: 51223

2.2 初始化导入的图片

def load(image_file):
    image = tf.io.read_file(image_file)
    image = tf.image.decode_jpeg(image)
    image = tf.cast(image, tf.float32)
    image = tf.image.resize(image, [image_size, image_size])
    image = (image - 127.5) / 127.5

    return image

2.3 构建数据集对象

dataset = X_train.map(
    load, num_parallel_calls=tf.data.experimental.AUTOTUNE).cache().shuffle(
    SHUFFLE_SIZE).batch(batch_size).repeat(100)

3、构建DCGAN网络

3.1 建立生成器

class Generator(tf.keras.Model):
    def __init__(self):
        super(Generator, self).__init__()
        filter = 64
        # 转置卷积层1,输出channel为filter*8,核大小4,步长1,不使用padding,不使用偏置
        self.conv1 = layers.Conv2DTranspose(filter*8, 4,1, 'valid', use_bias=False)
        self.bn1 = layers.BatchNormalization()
        # 转置卷积层2
        self.conv2 = layers.Conv2DTranspose(filter*4, 4,2, 'same', use_bias=False)
        self.bn2 = layers.BatchNormalization()
        # 转置卷积层3
        self.conv3 = layers.Conv2DTranspose(filter*2, 4,2, 'same', use_bias=False)
        self.bn3 = layers.BatchNormalization()
        # 转置卷积层4
        self.conv4 = layers.Conv2DTranspose(filter*1, 4,2, 'same', use_bias=False)
        self.bn4 = layers.BatchNormalization()
        # 转置卷积层5
        self.conv5 = layers.Conv2DTranspose(3, 4,2, 'same', use_bias=False)

    def call(self, inputs, training=None):
        x = inputs # [z, 100]
        # Reshape乘4D张量，方便后续转置卷积运算:(b, 1, 1, 100)
        x = tf.reshape(x, (x.shape[0], 1, 1, x.shape[1]))
        x = tf.nn.relu(x) # 激活函数
        # 转置卷积-BN-激活函数:(b, 4, 4, 512)
        x = tf.nn.relu(self.bn1(self.conv1(x), training=training))
        # 转置卷积-BN-激活函数:(b, 8, 8, 256)
        x = tf.nn.relu(self.bn2(self.conv2(x), training=training))
        # 转置卷积-BN-激活函数:(b, 16, 16, 128)
        x = tf.nn.relu(self.bn3(self.conv3(x), training=training))
        # 转置卷积-BN-激活函数:(b, 32, 32, 64)
        x = tf.nn.relu(self.bn4(self.conv4(x), training=training))
        # 转置卷积-激活函数:(b, 64, 64, 3)
        x = self.conv5(x)
        x = tf.tanh(x) # 输出x范围0~1,与预处理一致

        return x

3.2 建立判别器

class Discriminator(tf.keras.Model):
    # 判别器
    def __init__(self):
        super(Discriminator, self).__init__()
        filter = 64
        # 卷积层
        self.conv1 = layers.Conv2D(filter, 4, 2, 'valid', use_bias=False)
        self.bn1 = layers.BatchNormalization()
        # 卷积层
        self.conv2 = layers.Conv2D(filter*2, 4, 2, 'valid', use_bias=False)
        self.bn2 = layers.BatchNormalization()
        # 卷积层
        self.conv3 = layers.Conv2D(filter*4, 4, 2, 'valid', use_bias=False)
        self.bn3 = layers.BatchNormalization()
        # 卷积层
        self.conv4 = layers.Conv2D(filter*8, 3, 1, 'valid', use_bias=False)
        self.bn4 = layers.BatchNormalization()
        # 卷积层
        self.conv5 = layers.Conv2D(filter*16, 3, 1, 'valid', use_bias=False)
        self.bn5 = layers.BatchNormalization()
        # 全局池化层
        self.pool = layers.GlobalAveragePooling2D()
        # 特征打平
        self.flatten = layers.Flatten()
        # 2分类全连接层
        self.fc = layers.Dense(1)


    def call(self, inputs, training=None):
        # 卷积-BN-激活函数:(4, 31, 31, 64)
        x = tf.nn.leaky_relu(self.bn1(self.conv1(inputs), training=training))
        # 卷积-BN-激活函数:(4, 14, 14, 128)
        x = tf.nn.leaky_relu(self.bn2(self.conv2(x), training=training))
        # 卷积-BN-激活函数:(4, 6, 6, 256)
        x = tf.nn.leaky_relu(self.bn3(self.conv3(x), training=training))
        # 卷积-BN-激活函数:(4, 4, 4, 512)
        x = tf.nn.leaky_relu(self.bn4(self.conv4(x), training=training))
        # 卷积-BN-激活函数:(4, 2, 2, 1024)
        x = tf.nn.leaky_relu(self.bn5(self.conv5(x), training=training))
        # 卷积-BN-激活函数:(4, 1024)
        x = self.pool(x)
        # 打平
        x = self.flatten(x)
        # 输出，[b, 1024] => [b, 1]
        logits = self.fc(x)

        return logits

3.3 实例化判别器和生成器

generator = Generator() # 创建生成器
generator.build(input_shape = (4, z_dim))
discriminator = Discriminator() # 创建判别器
discriminator.build(input_shape=(4, 64, 64, 3))

其中input_shape 的第一个维度其实随便设一个数就可以，但是用None 会报错，不知道为什么……

4、建立损失函数

4.1 生成器损失函数

loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=True)
def g_loss_fn(d_fake_logits):
    # 计算生成图片与1之间的误差
    loss = tf.reduce_mean(loss_object(tf.ones_like(d_fake_logits), d_fake_logits))

    return loss

4.2 判别器损失函数

def d_loss_fn(d_real_logits, d_fake_logits):
    # 真实图片与1之间的误差
    d_loss_real = tf.reduce_mean(loss_object(tf.ones_like(d_real_logits), d_real_logits))
    # 生成图片与0之间的误差
    d_loss_fake = tf.reduce_mean(loss_object(tf.zeros_like(d_fake_logits), d_fake_logits))
    # 合并误差
    loss = d_loss_fake + d_loss_real

    return loss

5、初始化优化器

# 分别为生成器和判别器创建优化器
learning_rate = 0.0002
g_optimizer = keras.optimizers.Adam(learning_rate=learning_rate, beta_1=0.5)
d_optimizer = keras.optimizers.Adam(learning_rate=learning_rate, beta_1=0.5)

6、定义梯度下降过程

def train_step(batch_x):
    # 采样隐藏向量
    batch_z = tf.random.normal([batch_size, z_dim])
    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        # 采样生成图片
        fake_image = generator(batch_z, training=True)
        # 判定生成图片
        d_fake_logits = discriminator(fake_image, training=True)
        # 判定真实图片
        d_real_logits = discriminator(batch_x, training=True)
        d_loss = d_loss_fn(d_real_logits, d_fake_logits)
        g_loss = g_loss_fn(d_fake_logits)
    grads_d = disc_tape.gradient(d_loss, discriminator.trainable_variables)
    grads_g = gen_tape.gradient(g_loss, generator.trainable_variables)
    d_optimizer.apply_gradients(zip(grads_d, discriminator.trainable_variables))
    g_optimizer.apply_gradients(zip(grads_g, generator.trainable_variables))
    
    return d_loss, g_loss

注意，要将所有关于生成器、判别器的操作都放在梯度带中。

7、将生成的多张图像放到一个图里并保存

def save_result(val_out, val_block_size, image_path, color_mode):
    preprocesed = ((val_out + 1.0) * 127.5).astype(np.uint8)
    final_image = np.array([])
    single_row = np.array([])
    for b in range(val_out.shape[0]):
        # concat image into a row
        if single_row.size == 0:
            single_row = preprocesed[b, :, :, :]
        else:
            single_row = np.concatenate((single_row, preprocesed[b, :, :, :]), axis=1)

        # concat image row to final_image
        if (b+1) % val_block_size == 0:
            if final_image.size == 0:
                final_image = single_row
            else:
                final_image = np.concatenate((final_image, single_row), axis=0)

            # reset single row
            single_row = np.array([])
    Image.fromarray(final_image).save(image_path)

8、训练

for n, data in dataset.enumerate():
    d_loss, g_loss = train_step(data)
    print('.', end='')
    if n % 100 == 0:
        print()
        print(n.numpy(), 'd-loss:',float(d_loss), 'g-loss:', float(g_loss))
        # 可视化
        z = tf.random.normal([100, z_dim])
        fake_image = generator(z, training=False)
        img_path = os.path.join('gan_images', 'gan-%d.png'%n)
        save_result(fake_image.numpy(), 10, img_path, color_mode='P')