GAN 源码学习

GAN 源码学习

学习一下 github 上面用 Tensorflow 写的 GAN 代码。

源代码链接：https://github.com/huzixuan1/TF_2.0/blob/master/GAN

dataset.py

import multiprocessing # 允许程序员充分利用多个核心
import tensorflow as tf # 导入 tensorflow

def make_anime_dataset(img_paths,batch_size,resize=64,drop_remainder=True,shuffle=True,repeat=1):  
    # 生成动漫数据集（图片路径, 批量大小, ...）
    @tf.function

Tip: 函数装饰器

# 用函数 function1 去装饰函数 function2

def function1(function): # 函数 function1 的形式参数传入一个函数
    # ...
    function()
    return A

@function1
def function2(): 
    # ...

# 上面的代码等价于下面的代码
function2=function1(function2)
# 相当于用函数 function1 处理了一下 function2
    

Tip: tf.function
TensorFlow 2 中，只需要将我们希望以图执行模式（Graph Execution）/ 而非及时执行模式（Eager Execution）的代码写成一个函数，并在函数前加上@tf.function 进行装饰就可以了。

def make_anime_dataset(img_paths, batch_size, resize=64, drop_remainder=True, shuffle=True, repeat=1):  # img_paths 是一个数组
    @tf.function # 下面函数内的代码以图执行模式执行
    def _map_fn(img):
        img = tf.image.resize(img, [resize, resize]) # 将图片大小变为 64x64
        img = tf.clip_by_value(img,0,255) # 将张量 img 中的数值限制在 0~255 之内
        img = img / 127.5 - 1 # 将图片的值归一化为 -1 ~ 1（猜测后面的会用 tanh 作为激活函数）
        return img # 返回归一化处理后的图片张量
    
    dataset = disk_imag_batch_dataset(img_paths, 
                                      batch_size, 
                                      drop_remainder=drop_remainder,
                                      map_fn=map_fn,
                                      shuffle=shuffle,
                                      repeat=repeat)
    img_shape = (resize, resize, 3) # 图片的张量形状 （64x64x3） 3 个通道为 RGB 三个通道
    len_dataset = len(img_paths) // batch_size  # 有多少个 batch // 表示整除
    
    return dataset, img_shape, len_dataset

def batch_dataset(dataset, 
                  batch_size, 
                  drop_remainder=True, 
                  n_prefetch_batch=1,
                  flter_fn=None,
                  n_map_threads=None,
                  filter_after_map=False,
                  shuffle=True,
                  shuffle_buffer_size=None,
                  repeat=None):
    # 对数据集进行 洗牌、筛选、处理、打包
    if n_map_threads is None: # 如果参数 n_map_threads 传入值为 None
        n_map_threads = multiprocessing.cpu_count() # 返回 cpu 的核数 目前 cpu 应该都是 4 核的吧，老一点的 cpu 有 2 核的
    if shuffle and shuffle_buffer_size is None: # 如果这两个参数其中一个是 None
        shuffle_buffer_size = max(batch_size*128,2048) # 最小的用于打乱的 buffer_size 是2048，其他的是批量大小的 128 倍
    
    # 在对 dataset 执行 map 之前，先执行 shuffle 是更有效率的，因为 map 有时很耗时间
    if shuffle： # 如果需要打乱
        dataset=dataset.shuffle(shuffle_buffer_size) # 对数据集进行打乱
    
    if not filter_after_map:
        if filter_fn:
            dataset = dataset.filter(filter_fn) # 对数据集进行过滤，过滤标准依照函数filter_fn
            
        if map_fn:
            dataset = dataset.map(map_fn, num_parallel_calls=n_map_threads) # 对数据集中的所有数据进行操作，多线程并行操作       
    else:  # 如果不需要过滤
        if map_fn:
            dataset = dataset.map(map_fn, num_parallel_calls=n_map_threads)
        
        if filter_fn:
            dataset = dataset.filter(filter_fn)
    dataset = dataset.batch(batch_size, drop_remainder=drop_remainder) # 将数据集分成一个一个小块, drop_remainder:对于最后一个不足 batch_size 的数据是保留还是删除,默认选择保留。
    dataset = dataset.repeat(repeat).prefetch(n_prefetch_batch) # prefetch 放在最后提供一个 software pipelining 机制。使得 GPU 在训练上一次准备的数据时，CPU 去准备下一次需要的数据。如果单个训练步骤消耗 n 个元素，则添加 prefetch(n)。
    return dataset

def memory_data_batch_dataset(memory_data, #list/ndarray/Tensor
                              batch_size,
                             drop_remainder=True,
                             n_prefectch_batch=1,
                             filter_fn=None,
                             map_fn=None,
                             n_map_threads=None,
                             filter_after_map=False,
                             shuffle_buffer_size=None,
                             repeat=None):
    # 从内存中读取数据、制作数据集
    dataset = tf.data.Dataset.from_tensor_slices(memory_data) # 将 memory_data 制作成数据集
    dataset = batch_dataset(dataset,
                            batch_size，
                           drop_remainder=drop_remainder,
                           n_prefetch_batch=n_prefetch_batch,
                           filter_fn=filter_fn,
                           map_fn=map_fn,
                           n_map_threads=n_map_threads,
                           filter_after_map=filter_after_map,
                           shuffle=shuffle,
                           shuffle_buffer_size=shuffle_buffer_size,
                           repeat=repeat) # 对 dataset 洗牌、筛选、处理、打包
    return dataset

def disk_image_batch_dataset(img_paths,
                            batch_size,
                            labels=None,
                            drop_remainder=True,
                            n_prefetch_batch=1,
                            filter_fn=None,
                            map_fn=None,
                            n_map_threads=None,
                            filter_after_map=False,
                            shuffle=True,
                            shuffle=buffer_size=None,
                            repeat=None):
    # 从磁盘数据中读取数据制作数据集
    if label is None:
        memory_data = img_paths # 如果没有标签
    else:
        memory_data = (img_patchs,labels) # 如果有标签
    
    def parse_fn(path, *label): # * 表示接受一个元组、** 表示接受一个字典
        # 读取文件 返回一个元组
        img = tf.io.read_file(path) # 读取文件
        img = tf.image.decode_png(img,3) # 解码 png 图片 通道数固定到 RGB 三个通道
        return (img,) + label # 合成一个元组
    
    if map_fn: # 融合 map_fn 和 parse_fn 如果 map_fn 有传入值
        def map_fn_(*args):
            return map_fn(*parse_fn(*args)) # 对图片进行处理
    else:
        map_fn_ = parse_fn # 将 map_fn 和 parse_fn合成为一个函数 map_fn_
    
    # 从磁盘中读取数据，其实本质上也是把磁盘中的数据读取到内存中，最后从内存中制作数据集
    dataset = memory_data_batch_dataset(memory_data,
                                       batch_size,
                                       drop_remainder=drop_remainder,
                                       n_prefetch_batch=n_prefetch_batch,
                                       filter_fn=filter_fn,
                                       map_fn=map_fn_,
                                       n_map_threads=n_map_threads,
                                       filter_after_map=filter_after_map,
                                       shuffle=shuffle,
                                       shuffle_buffer_size=shuffle_buffer_size,
                                       repeat=repeat)
    return dataset

gan.py

import os # 导入 os 系统模块
import tensorflow as tf # 导入 tensorflow 模块
from tensorflow import keras # 导入 keras 模型部分

os.environp['TF_CPP_MIN_LOG_LEVEL']='2' # 屏蔽 INFO 和 WARNING 输出 ERRPR 和 FATAL
class Generator(keras.Model):
    def __init__(self):
        # 显示确定一些需要的参数、其实可以不用在__init__中定义，直接在 call 中使用即可。
        super(Generator,self).__init__() # 向父类注册
        self.fc=keras.layers.Dense(3*3*512)
        
        self.conv1=keras.layers.Conv2DTranspose(256,3,3,'valid') # 反卷积 kernel_size=3 filters=256,padding='valid'
        self.bn1=keras.layers.BatchNormalization()
        self.conv2=keras.layers.Conv2DTranspose(128,5,2,'valid')
        self.b2=keras.layers.BatchNormalization()
        
        self.conv3=keras.layers.Conv2DTranspose(3,4,3,'valid')
    
    def call(self, inputs, training=None, mask=None):
        x=self.fc(inputs)
        x=tf.reshape(x,[-1,3,3,512])
        x=tf.nn.leaky_relu(x)
        
        x=self.conv1(x)
        x=self.bn1(x,training=training)
        x=tf.nn.leaky_relu(x)
       
        x=self.conv2(x)
        x=self.bn2(x,training=training)
        x=tf.nn.leaky_relu(x)
        
        x=self.conv3(x)
        x=tf.tanh(x)
        return x
        # 这个例子告诉我们，在 keras 的自定义层中直接在 call 中调用 keras 内置层是可以有可训练变量的。
        
 class Discriminator(keras.Model): 
    def __init__(self):
        super(Discriminator,self).__init__()
        self.conv1=keras.layers.Conv2D(64,5,3,'valid')
        
        self.conv2=keras.layers.Conv2D(128,5,3,'valid')
        self.bn2=keras.layers.BatchNormalization()
        
        self.conv3=keras.layers.BatchNormalization()
        self.bn3=keras.layers.BatchNormalization()
        
        self.flattn=keras.layers.Flatten()
        self.fc=keras.layers.Dense(1)
        
    def call(self, inputs, training=None, mask=None):
        x=self.conv1(inputs)
        x=tf.nn.leaky_relu(x)
        x=self.conv2(x)
        x=self.bn2(x,training=training)
        x=tf.nn.leaky_relu(x)
        x=self.conv3(x)
        x=self.bn3(x,training=training)
        x=tf.nn.leaky_relu(x)
        
        x=self.flatten(x)
        logits=self.fc(x)
        
        return logits

Tip: keras 两种初始化模型的方法

原来我只会第一种初始化模型的方法，这份源码教会了我第二种方法

# 方法1：从 Input 开始指定前向过程，最后根据输入和输出来建立模型
inputs = tf.keras.Input(shape=(3,))
x = tf.keras.layers.Dense(4, activation=tf.nn.relu)(inputs)
outputs = tf.keras.layers.Dense(5, activation=tf.nn.softmax)(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)

# 方法2：构建 Model 的子类，__init__中定义层的实现，call函数中实现前向过程

class MyModel(tf.keras.Model):

    def __init__(self):
        super(MyModel, self).__init__()
        self.dense1 = tf.keras.layers.Dense(4, activation=tf.nn.relu)
        self.dense2 = tf.keras.layers.Dense(5, activation=tf.nn.softmax)
        
    def call(self, inputs):
        x = self.dense1(inputs)
        output = self.dense2(2)
        return output
    
model = MyModel()    # 通过初始化整个类来初始化该模型

def main():
    d=Discrimination()
    g=Generator()
    
    x=tf.random.normal([2,64,64,3])
    z=tf.random.normal([2,100]) # 随机初始化两个一百维度的 z
    
    prob=d(x)
    print(prob.shape) # 应该的形状为[-1,1]
    x_hat=g(z)
    print(x_hat.shape) # 应该的形状为[-1,64,64,3]
if __name__=='__main__': # 这样写 在作为一个模块导入其他文件时 main() 并不会被直接执行因为在其他文件调用此文件时，__name__!='__main__',而如果直接运行本模块的化，可以直接运行 main 函数此时 __name__=='main'
    main()

gan_train.py

import os
import numpy as np
import tensorflow as tf
from tensorflow import keras
form PIL import Image # PIL 是图像处理标准库
import glob # glob 是操作文件的相关模块
from gan import Generator, Discriminator

from dataset import make_anime_dataset

def save_result(val_out, val_block_size, image_path, color_mode):
    # 保存结果
    def prepocess(img):
        img = ((img+1.0)*127.5).astype(np.uint8) # 将输出的张量复原为一张图像
        return img
    
    preprocesed = prepocess(val_out)
    final_image = np.array([])
    single_row = np.array([])
    for b in range(val_out.shape[0]):
        if single_row.size == 0: # 如果 single_row 中没有图片
            single_row = preprocesed[b,:,:,:] # 取第 b 张图片
        else: # 如果 single_row 中已经有图片了
            single_row = np.concatenate((single_row,preprocesed[b,:,:,:]),axis=1) # 将图片和之前的图片沿着 axis=1轴 合并成一张图片
        
        if (b+1) % val_block_size == 0: # 如果这是最后一张图片
            if final_image.size == 0:
                final_image = single_row # 将 final_image 令为最后一张图片
            else:
                final_image = np.concatenate((final_image, single_row), axis=0) # 如果 final_image 中有图片就再在其中添加 single_row
            
            single_row = np.array([]) # 重置 single_row
    
    if final_image.shape[2]==1:
        final_image=np.squeeze(final_image,axis=2)
    Image.fromarray(final_image).save(image_path)
    

def celoss_ones(logits): # logits shape (None,1)
    # 与全 1 张量之间的交叉熵
    loss = tf.nn.sigmoid_cross_entropy_with_logits(logits=logits,
                                                   labels=tf.ones_like(logits)) # tf.zeros_like/tf.ones_like 新建一个与给定 tensor 类型大小一致的 tensor 所有元素为 0/1 .
    return tf.reduce_mean(loss)


def celoss_zeros(logits): # logits shape (None,1)
    # 与全 0 张量之间的交叉熵
    loss = tf.nn.sigmoid_cross_entropy_with_logits(logits=logits,
                                                   labels=tf.zeros_like(logits)) # tf.zeros_like/tf.ones_like 新建一个与给定 tensor 类型大小一致的 tensor 所有元素为 0/1 .
    return tf.reduce_mean(loss)


def d_loss_fn(generator, discriminator, batch_z, batch_x, is_training):    		     
    # discriminator 的损失函数
    fake_image = generator(batch_z, is_training)
    d_fake_logits = discriminator(fake_image, is_training)
    d_real_logits = discriminator(batch_x, is_training)
    
    d_loss_real = celoss_ones(d_real_logits)
    d_loss_fake = celoss_zeros(d_fake_logits)
    
    loss = d_loss_fake + d_loss_real # 注意 fake 和 real 是等比例输入的，之后 discriminator 的 loss 应该将两部分加起来。
    
    return loss

def g_loss_fn(generator, discriminator, batch_z, is_training):
    # generator 的损失函数
    fake_image = generator(batch_z, istraining) # 相当于一个 layer (batch_z 是该 layer 的输入)
    d_fake_logits = discriminator(fake_image,is_training)
    loss = celoss_ones(d_fake_logits)
    
    return loss

def main():
    
    tf.random.set_seed(22) # 设置随机变量种子
    np.random.seed(22)
    os.environ['TF_CPP_MIN_LOG_LEVEL']='2' # ERROR 及其以上的错误才报
    assert tf.__version__.startswith('2.') # 一定要是 tensorflow 2.0
    
    

Tip: assert

assert expression # 声称，如果不是的话就抛出错误

# 等价于：

if not expression:
    raise AssertionError
    

def main():
    tf.random.set_seed(22) # 设置随机变量种子
    np.random.seed(22)
    os.environ['TF_CPP_MIN_LOG_LEVEL']='2' # ERROR 及其以上的错误才报
    assert tf.__version__.startswith('2.') # 一定要是 tensorflow 2.
    
    z_dim = 100
    epochs = 3000000
    batch_size = 512
    learning_rate = 0.002
    is_training = True
    
    img_path=glob.glob(r'E:\python_pro\TF2.0\GAN\faces\*.jpg') #图片路径
    dataset, img_shape, _ = make_anime_dataset(img_path, batch_size)
    print(dataset, img_shape)
    dataset = dataset.repeat()
    db_iter = iter(dataset) # iter(object)函数用来生成迭代器，其中 object 为支持迭代的集合对象，返回一个迭代器对象
    
    generator = Generator()
    generator.build(input_shape=(None, z_dim))
    discriminator = Discriminator()
    discriminator.build(input_shape=(None, 64, 64, 3))
    
    g_optimizer = tf.optimizers.Adam(learning_rate=learning_rate, beta_1=0.5)
    d_optimizer = tf.optimizers.Adam(learning_rate=learning_rate, beta_1=0.5)
    
    for epoch in range(epochs): # 每一个训练周期
        
        batch_z = tf.random.uniform([batch_size,z_dim],minval=-1.,maxval=1.)
        batch_x = next(db_iter) # 每个周期只会取其中的一个小 batch 来训练

Tip:以前我每个 epoch 选取 batch 的写法

for epoch in range(epochs):
    for data, label in dataset:

没有什么太大的区别，它的一个周期取一个 batch 我的一个周期 遍历整个数据集。

# 发现一个新大陆，所有的Dataset 可以用 iter(Dataset) 生成迭代对象，使用 next 取出。相当于一个栈式结构，每次取出一个 tf.Tensor.
db_iter=iter(dataset)
for epoch in range(epochs):
    batch_x=next(db_iter)

        with if.GradientTape() as tape:
            d_loss = d_loss_fn(generator, discriminator, batch_z, batch_x, is_training)
            grads = tape.gradient(d_loss,discriminator.trainable_variables)  
            d_optimizer.apply_gradients(zip(grads, discriminator.trainable_variables))
            
        with tf.GradientTape() as tape:
            g_loss = g_loss_fn(generator, discriminator, batch_z, is_training)
            grads = tape.gradient(g_loss,generator.trainable_variables)
            g_optimizer.apply_gradients(zip(grads, generator.trainable_variables))
        
        if epoch % 100 == 0:
            print(epoch,'d-loss:',float(d_loss),'g-loss',float(g_loss))
            
            z = tf.random.uniform([100,z_dim])
            fake_image = generator(z,training=False)
            img_path = os.path.join('images','gan-%d.png'%epoch)
            save_result(fake_image.numpy(),10,img_path,color_mode='P')



if __name__=='__main__':
    main()
    

学学别人写的源码也是挺好的。