最近对对抗生成网络GAN比较感兴趣,相关知识点文章还在编辑中,以下这个是一个练手的小项目~
(在原模型上做了,为了减少计算量让其好训练一些。)
一、导入工具包
-
import tensorflow
as tf
-
from tensorflow.keras
import layers
-
-
import numpy
as np
-
import os
-
import time
-
import glob
-
import matplotlib.pyplot
as plt
-
from IPython.display
import clear_output
-
from IPython
import display
1.1 设置GPU
-
gpus = tf.config.list_physical_devices(
"GPU")
-
-
if gpus:
-
gpu0 = gpus[
0]
#如果有多个GPU,仅使用第0个GPU
-
tf.config.experimental.set_memory_growth(gpu0,
True)
#设置GPU显存用量按需使用
-
tf.config.set_visible_devices([gpu0],
"GPU")
-
gpus
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
二、导入训练数据
链接: 点这里
-
fileList = glob.glob(
'./ani_face/*.jpg')
-
len(fileList)
41621
2.1 数据可视化
-
# 随机显示几张图
-
for index,i
in
enumerate(fileList[:
3]):
-
display.display(display.Image(fileList[index]))
2.2 数据预处理
-
# 文件名列表
-
path_ds = tf.data.Dataset.from_tensor_slices(fileList)
-
-
# 预处理,归一化,缩放
-
def
load_and_preprocess_image(
path):
-
image = tf.io.read_file(path)
-
image = tf.image.decode_jpeg(image, channels=
3)
-
image = tf.image.resize(image, [
64,
64])
-
image /=
255.0
# normalize to [0,1] range
-
image = tf.reshape(image, [
1,
64,
64,
3])
-
return image
-
-
image_ds = path_ds.
map(load_and_preprocess_image)
-
image_ds
<MapDataset shapes: (1, 64, 64, 3), types: tf.float32>
-
# 查看一张图片
-
for x
in image_ds:
-
plt.axis(
"off")
-
plt.imshow((x.numpy() *
255).astype(
"int32")[
0])
-
break
三、网络构建
3.1 D网络
-
discriminator = keras.Sequential(
-
[
-
keras.Input(shape=(
64,
64,
3)),
-
layers.Conv2D(
64, kernel_size=
4, strides=
2, padding=
"same"),
-
layers.LeakyReLU(alpha=
0.2),
-
layers.Conv2D(
128, kernel_size=
4, strides=
2, padding=
"same"),
-
layers.LeakyReLU(alpha=
0.2),
-
layers.Conv2D(
128, kernel_size=
4, strides=
2, padding=
"same"),
-
layers.LeakyReLU(alpha=
0.2),
-
layers.Flatten(),
-
layers.Dropout(
0.2),
-
layers.Dense(
1, activation=
"sigmoid"),
-
],
-
name=
"discriminator",
-
)
-
discriminator.summary()
Model: "discriminator" _________________________________________________________________ Layer (type) Output Shape Param # ================================================================= conv2d (Conv2D) (None, 32, 32, 64) 3136 _________________________________________________________________ leaky_re_lu (LeakyReLU) (None, 32, 32, 64) 0 _________________________________________________________________ conv2d_1 (Conv2D) (None, 16, 16, 128) 131200 _________________________________________________________________ leaky_re_lu_1 (LeakyReLU) (None, 16, 16, 128) 0 _________________________________________________________________ conv2d_2 (Conv2D) (None, 8, 8, 128) 262272 _________________________________________________________________ leaky_re_lu_2 (LeakyReLU) (None, 8, 8, 128) 0 _________________________________________________________________ flatten (Flatten) (None, 8192) 0 _________________________________________________________________ dropout (Dropout) (None, 8192) 0 _________________________________________________________________ dense (Dense) (None, 1) 8193 ================================================================= Total params: 404,801 Trainable params: 404,801 Non-trainable params: 0
3.2 G网络
-
latent_dim =
128
-
-
generator = keras.Sequential(
-
[
-
keras.Input(shape=(latent_dim,)),
-
layers.Dense(
8 *
8 *
128),
-
layers.Reshape((
8,
8,
128)),
-
layers.Conv2DTranspose(
128, kernel_size=
4, strides=
2, padding=
"same"),
-
layers.LeakyReLU(alpha=
0.2),
-
layers.Conv2DTranspose(
256, kernel_size=
4, strides=
2, padding=
"same"),
-
layers.LeakyReLU(alpha=
0.2),
-
layers.Conv2DTranspose(
512, kernel_size=
4, strides=
2, padding=
"same"),
-
layers.LeakyReLU(alpha=
0.2),
-
layers.Conv2D(
3, kernel_size=
5, padding=
"same", activation=
"sigmoid"),
-
],
-
name=
"generator",
-
)
-
generator.summary()
3.3 重写 train_step
-
class
GAN(keras.Model):
-
def
__init__(
self, discriminator, generator, latent_dim):
-
super(GAN, self).__init__()
-
self.discriminator = discriminator
-
self.generator = generator
-
self.latent_dim = latent_dim
-
-
def
compile(
self, d_optimizer, g_optimizer, loss_fn):
-
super(GAN, self).
compile()
-
self.d_optimizer = d_optimizer
-
self.g_optimizer = g_optimizer
-
self.loss_fn = loss_fn
-
self.d_loss_metric = keras.metrics.Mean(name=
"d_loss")
-
self.g_loss_metric = keras.metrics.Mean(name=
"g_loss")
-
-
@property
-
def
metrics(
self):
-
return [self.d_loss_metric, self.g_loss_metric]
-
-
def
train_step(
self, real_images):
-
# 生成噪音
-
batch_size = tf.shape(real_images)[
0]
-
random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
-
-
# 生成的图片
-
generated_images = self.generator(random_latent_vectors)
-
-
# Combine them with real images
-
combined_images = tf.concat([generated_images, real_images], axis=
0)
-
-
# Assemble labels discriminating real from fake images
-
labels = tf.concat(
-
[tf.ones((batch_size,
1)), tf.zeros((batch_size,
1))], axis=
0
-
)
-
# Add random noise to the labels - important trick!
-
labels +=
0.05 * tf.random.uniform(tf.shape(labels))
-
-
# 训练判别器,生成的当成0,真实的当成1
-
with tf.GradientTape()
as tape:
-
predictions = self.discriminator(combined_images)
-
d_loss = self.loss_fn(labels, predictions)
-
grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
-
self.d_optimizer.apply_gradients(
-
zip(grads, self.discriminator.trainable_weights)
-
)
-
-
# Sample random points in the latent space
-
random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
-
-
# Assemble labels that say "all real images"
-
misleading_labels = tf.zeros((batch_size,
1))
-
-
# Train the generator (note that we should *not* update the weights
-
# of the discriminator)!
-
with tf.GradientTape()
as tape:
-
predictions = self.discriminator(self.generator(random_latent_vectors))
-
g_loss = self.loss_fn(misleading_labels, predictions)
-
grads = tape.gradient(g_loss, self.generator.trainable_weights)
-
self.g_optimizer.apply_gradients(
zip(grads, self.generator.trainable_weights))
-
-
# Update metrics
-
self.d_loss_metric.update_state(d_loss)
-
self.g_loss_metric.update_state(g_loss)
-
return {
-
"d_loss": self.d_loss_metric.result(),
-
"g_loss": self.g_loss_metric.result(),
-
}
3.4 设置回调函数
-
class
GANMonitor(keras.callbacks.Callback):
-
def
__init__(
self, num_img=3, latent_dim=128):
-
self.num_img = num_img
-
self.latent_dim = latent_dim
-
-
def
on_epoch_end(
self, epoch, logs=None):
-
random_latent_vectors = tf.random.normal(shape=(self.num_img, self.latent_dim))
-
generated_images = self.model.generator(random_latent_vectors)
-
generated_images *=
255
-
generated_images.numpy()
-
for i
in
range(self.num_img):
-
img = keras.preprocessing.image.array_to_img(generated_images[i])
-
display.display(img)
-
img.save(
"gen_ani/generated_img_%03d_%d.png" % (epoch, i))
四、训练模型
-
epochs =
100
# In practice, use ~100 epochs
-
-
gan = GAN(discriminator=discriminator, generator=generator, latent_dim=latent_dim)
-
gan.
compile(
-
d_optimizer=keras.optimizers.Adam(learning_rate=
0.0001),
-
g_optimizer=keras.optimizers.Adam(learning_rate=
0.0001),
-
loss_fn=keras.losses.BinaryCrossentropy(),
-
)
-
-
gan.fit(
-
image_ds, epochs=epochs, callbacks=[GANMonitor(num_img=
10, latent_dim=latent_dim)]
-
)
五、保存模型
-
#保存模型
-
gan.generator.save(
'./data/ani_G_model')
生成模型文件:点这里
六、生成漫画脸
-
G_model = tf.keras.models.load_model(
'./data/ani_G_model/',
compile=
False)
-
-
def
randomGenerate():
-
noise_seed = tf.random.normal([
16,
128])
-
predictions = G_model(noise_seed, training=
False)
-
fig = plt.figure(figsize=(
8,
8))
-
for i
in
range(predictions.shape[
0]):
-
plt.subplot(
4,
4, i+
1)
-
img = (predictions[i].numpy() *
255 ).astype(
'int')
-
plt.imshow(img )
-
plt.axis(
'off')
-
plt.show()
-
count =
0
-
while
True:
-
randomGenerate()
-
clear_output(wait=
True)
-
time.sleep(
0.1)
-
if count >
100:
-
break
-
count+=
1
转载:https://blog.csdn.net/suic009/article/details/128120410
查看评论