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.functionTip: 函数装饰器
# 用函数 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 datasetgan.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 logitsTip: 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()
学学别人写的源码也是挺好的。