当前位置:   article > 正文

深度学习(6):图像超分辨率(Image Super Resolution)重建_图像超分辨率重建

图像超分辨率重建

目标:使用深度学习技术对图像进行超分辨率重建,涉及到的技术包括卷积神经网络,生成对抗网络,残差网络等。

技术:本实验使用到“Jupyter Notebook”等开发组件,涉及到了“TensorFlow”、“NumPy”、“scipy. misc”、“PIL.image”等框架和库,其中“scipy. misc”与“PIL.image”用于图像处理。

数据集

本实验的数据可以选择CV领域各类常见的数据集,实验将以CelebA数据集作为示例。CelebA是香港中文大学开放的人脸识别数据,包括了10177个名人的202599张图片,并有5个位置标记和40个属性标记,可以作为人脸检测、人脸属性识别、人脸位置定位等任务的数据集。数据集可在Google Drive中下载,详细信息可在官方网站中查看。本实验使用该数据集中img_align_celeba.zip这一文件,选择了其中前10661张图片,每张图片根据人像双眼的位置调整为了219x178的尺寸。

本实验需要得到图像的低分辨率图像,并通过深度学习将这些低分辨率图像提升到高分辨率图像,最后与原图进行对比查看模型的效果。

在图像超分辨率问题中,理论上可以选择任意的图像数据,当时根据经验,使用有更多细节纹理的图像会有更好的效果,使用无损压缩格式的PNG格式图像比JPG格式图像有更好的效果。

一、原理介绍

1.GAN

可以参考博客GANs:生成对抗网络系列及应用_studyeboy的博客-CSDN博客_gans的应用

GAN解决了非监督学习中的著名问题:给定一批样本,训练一个系统,能够生成类似的新样本。

 

 

GAN的优势:

  • 根据实际的结果,看上去产生了更好的样本;
  • GAN能训练任何一种生成器网络;
  • GAN不需要设计遵循任何种类的因式分解的模型,任何生成器网络和任何鉴别器都会有用;
  • GAN无需利用马尔科夫链反复采样,无需在学习过程中进行推断,回避了近似计算棘手的概率的难题。

 

 

GAN的应用:

  • DCGAN应用到人脸图像
  • 计算机视觉,重建图像
  • 图像超分辨率
  • 图像去雨
  • 二次元头像炼成
  • 图片生成
  • 图片风格转换

 2.Generation

生成(generation)就是模型通过学习一些数据,然后生成类似的数据。让机器看一些动物图片,然后自己来产生动物的图片,这就是生成。

作为一种无监督或者自监督算法,自编码器本质上是一种数据压缩算法。

  • AutoEncoder通常有两个方面的应用:
    1. 是数据去噪,
    2. 是为进行可视化而降维。
    3. 进行图像压缩
    4. 传统自编码器被用于降维或特征学习

 参考博客:深度学习-生成模型:AutoEncoder模型(Encoder-Decoder框架==Embedding+Generator)【Variational AutoEncoder已被GAN取代】_u013250861的博客-CSDN博客_autoencoder与gan

 

 VAE变分自动编码器结构是一个经典的autocoder模型,网络的组成也是编码器、解码器、loss。

 Encoder-Decoder是一个模型构架,是一类算法统称,并不是特指某一个具体的算法,在这个框架下可以使用不同的算法来解决不同的任务。首先,编码(encode)由一个编码器将输入序列转化成一个固定维度的稠密向量,解码(decode)阶段将这个激活状态生成目标译文。参考Encoder-Decoder综述理解(推荐)_SmileAda的博客-CSDN博客_encoder-decoder

 EBGAN的判别器使用Encoder-Decoder的结构。

二、实验过程

1.安装库

!pip3 install imageio

2.导入库

  1. import os
  2. import tensorflow as tf
  3. import numpy as np
  4. import imageio
  5. import random
  6. from PIL import Image
  7. tf.enable_eager_execution()

3.下载数据 

  1. import oss2
  2. access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', 'LTAI4G1MuHTUeNrKdQEPnbph')
  3. access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', 'm1ILSoVqcPUxFFDqer4tKDxDkoP1ji')
  4. bucket_name = os.getenv('OSS_TEST_BUCKET', 'mldemo')
  5. endpoint = os.getenv('OSS_TEST_ENDPOINT', 'https://oss-cn-shanghai.aliyuncs.com')
  6. # 创建Bucket对象,所有Object相关的接口都可以通过Bucket对象来进行
  7. bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
  8. # 下载到本地文件
  9. bucket.get_object_to_file('data/c12/super_resolution_dataset.zip', 'super_resolution_dataset.zip')
  10. bucket.get_object_to_file('data/c12/super_resolution_imgs.zip', 'super_resolution_imgs.zip')

4.解压数据集

  1. !unzip -q -o super_resolution_dataset.zip
  2. !unzip -q -o super_resolution_imgs.zip
  3. !rm -rf __MACOSX

5.数据预处理

本实验的数据预处理需要将图1中的原始图像整理为神经网络的对应输入与输出,并对输入做数据增强。在预处理前,将最后五张图像移动到新的文件夹中作为测试图像,其余图像作为训练图像。

图1 

  • 图像尺寸调整

图中元素图像尺寸为219x178,为了提升实验效率和效果,首先将训练与测试图像调整到128x128的尺寸,注意实验没有直接使用“resize”函数,因为“resize”函数进行下采样会降低图像的分辨率,实验使用了“crop”函数在图像中间进行裁剪,并将最后裁剪后的图像持久化保存。

  • 载入图像

使用TensorFlow的Dataset API,该API对数据集进行了高级的封装,可以对数据进行批量的载入、预处理、批次读取、shuffle、prefetch等操作。注意prefetch操作有内存需要,内存不足可以不需要进行prefetch;由于本实验的后续的网络结构较深,因此对显存有相当高的要求,如果显存足够或不足可以适当进行调整batch大小。

  1. def delete_or_makedir(dir):
  2. if tf.gfile.Exists(dir):
  3. tf.gfile.DeleteRecursively(dir)
  4. tf.gfile.MakeDirs(dir)
  5. def resize_image(filename, hr_dir, resize_dir):
  6. if not filename.endswith('jpg'):
  7. return
  8. image = Image.open(os.path.join(hr_dir, filename))
  9. half_the_width = image.size[0] / 2
  10. half_the_height = image.size[1] / 2
  11. image = image.crop(
  12. (
  13. half_the_width - 64,
  14. half_the_height - 64,
  15. half_the_width + 64,
  16. half_the_height + 64
  17. )
  18. )
  19. file, _ = os.path.splitext(filename)
  20. image.save(os.path.join(resize_dir, file + '-resized.png'))
  21. def prepare_train_dirs(checkpoint_dir, train_log_dir, delete_train_log_dir=False):
  22. # Create checkpoint dir (do not delete anything)
  23. if not tf.gfile.Exists(checkpoint_dir):
  24. tf.gfile.MakeDirs(checkpoint_dir)
  25. # Cleanup train log dir
  26. if delete_train_log_dir:
  27. delete_or_makedir(train_log_dir)

6.图像预处理

在图像处理中,常常在图像输入网络前对图像进行数据增强。数据增强有两个主要目的,其一是通过对图像的随机变换增加训练数据,其二是通过随机的增强使得训练出的模型尽可能少地受到无关因素的影响,增加图像的泛化能力。本节首先在文件中读取前文裁剪后的图像;然后对训练图像进行随机的左右翻转;并在一定范围内随机调整图像的饱和度、亮度、对比度和色相;然后将读取的RGB值规范化到[-1, 1]区间;最后使用双三次插值的方法将图像下采样四倍到32*32尺寸。

  1. # 图像预处理
  2. def preprocess_image(image_path, training=False):
  3. image_size = 128
  4. k_downscale = 4
  5. downsampled_size = image_size // k_downscale
  6. image = tf.read_file(image_path)
  7. image = tf.image.decode_jpeg(image, channels=3)
  8. if training:
  9. # 在训练图像时,利用随机翻转对图像进行预处理来增加训练数据
  10. image = tf.image.random_flip_left_right(image)
  11. # 随机调整属性,使训练得到的模型尽可能小的受到无关因素的影响.
  12. image = tf.image.random_saturation(image, 0.95, 1.05) # 饱和度
  13. image = tf.image.random_brightness(image, 0.05) # 亮度
  14. image = tf.image.random_contrast(image, 0.95, 1.05) # 对比度
  15. image = tf.image.random_hue(image, 0.05) # 色相
  16. label = (tf.cast(image, tf.float32) - 127.5) / 127.5 # normalize to [-1,1] range
  17. feature = tf.image.resize_images(image, [downsampled_size, downsampled_size], tf.image.ResizeMethod.BICUBIC)
  18. feature = (tf.cast(feature, tf.float32) - 127.5) / 127.5 # normalize to [-1,1] range
  19. # if training:
  20. # feature = feature + tf.random.normal(feature.get_shape(), stddev=0.03)
  21. return feature, label
  1. # 加载数据
  2. def load_data(data_dir, training=False):
  3. filenames = tf.gfile.ListDirectory(data_dir)
  4. filenames = [os.path.join(data_dir, f) for f in filenames]
  5. random.shuffle(filenames)
  6. image_count = len(filenames)
  7. image_ds = tf.data.Dataset.from_tensor_slices(filenames)
  8. image_ds = image_ds.map(lambda image_path: preprocess_image(image_path, training=training))
  9. BATCH_SIZE = 2
  10. image_ds = image_ds.batch(BATCH_SIZE)
  11. # image_ds = image_ds.prefetch(buffer_size=400)
  12. return image_ds

7.持久化测试数据

在前序预处理后,实验将把测试集的特征和标签数据持久化到本地,以在后续的训练中与模型的输出做对比。

  1. # 保存特征数据
  2. def save_feature_label(train_log_dir, test_image_ds):
  3. feature_batch, label_batch = next(iter(test_image_ds))
  4. feature_dir = train_log_dir + '0_feature/'
  5. label_dir = train_log_dir + '0_label/'
  6. delete_or_makedir(feature_dir)
  7. delete_or_makedir(label_dir)
  8. for i, feature in enumerate(feature_batch):
  9. if i > 5:
  10. break
  11. imageio.imwrite(feature_dir + '{:02d}.png'.format(i), feature)
  12. for i, label in enumerate(label_batch):
  13. if i > 5:
  14. break
  15. imageio.imwrite(feature_dir + '{:02d}.png'.format(i), feature)

8.模型设计——残差网络

本实验将使用GAN、CNN和ResNet的组合构建超分辨率模型。GAN的生成器中使用到的残差块与上采样的PixelShuffle,然后是GAN中生成器与判别器,最后是模型的训练过程。

残差网络引入了残差块的设计,残差块的输入为x,正常的模型设计的输出是多层神经网络的输出F(x),残差块将输入的x与两层的输出F(x)相加的结果H(x)作为残差块的输出。训练的目标是使得残差F(x)=H(x)-x逼近于0,即H(x)与x尽可能的近似。随着网络层次的加深,这样的设计保证了在后续层次中网络的仍有可学习内容。

残差块中weight layer是有64个特征图输出、卷积核大小为3*3,Relu激活函数也改为了PRelu。

  1. # 残差块定义
  2. class _IdentityBlock(tf.keras.Model):
  3. def __init__(self, filter, stride, data_format):
  4. super(_IdentityBlock, self).__init__(name='')
  5. bn_axis = 1 if data_format == 'channels_first' else 3
  6. self.conv2a = tf.keras.layers.Conv2D(
  7. filter, (3, 3), strides=stride, data_format=data_format, padding='same', use_bias=False)
  8. # self.bn2a = tf.keras.layers.BatchNormalization(axis=bn_axis)
  9. self.prelu2a = tf.keras.layers.PReLU(shared_axes=[1, 2])
  10. self.conv2b = tf.keras.layers.Conv2D(
  11. filter, (3, 3), strides=stride, data_format=data_format, padding='same', use_bias=False)
  12. # self.bn2b = tf.keras.layers.BatchNormalization(axis=bn_axis)
  13. def call(self, input_tensor):
  14. x = self.conv2a(input_tensor)
  15. # x = self.bn2a(x)
  16. # x = tf.nn.leaky_relu(x)
  17. x = self.prelu2a(x)
  18. x = self.conv2b(x)
  19. # x = self.bn2b(x)
  20. x = x + input_tensor
  21. return x

9.上采样 PixelShuffler

目标是将32x32的低分辨率图像超分辨率到128x128。第一种方法是Conv2DTranspose,它用的是反卷积,即卷积操作的逆,该方法会造成非常明显的噪声像素;第二种上采样方法是TensorFlow UpSampling2D + Conv2D的方法,该方法是CNN中常见的max pooling的逆操作,该方法会损失较多的信息,实验效果不佳。

最终选择了PixelShuffle作为上采样的方法。其操作如图所示。输入为H*W的低分辨率图像,首先通过卷积操作得到r^2个特征图(r为上采样因子,即图像放大的倍数),其中特征图的大小与低分辨率图的大小一致,然后通过周期筛选(periodic shuffing)得到高分辨率的图像。

将卷积部分的操作放到了GAN的生成器中,下面的代码展示了如何在r^2个特征图上做周期帅选得到目标高分辨率图像输出。

  1. def pixelShuffler(inputs, scale=2):
  2. size = tf.shape(inputs)
  3. batch_size = size[0]
  4. h = size[1]
  5. w = size[2]
  6. c = inputs.get_shape().as_list()[-1]
  7. # Get the target channel size
  8. channel_target = c // (scale * scale)
  9. channel_factor = c // channel_target
  10. shape_1 = [batch_size, h, w, channel_factor // scale, channel_factor // scale]
  11. shape_2 = [batch_size, h * scale, w * scale, 1]
  12. # Reshape and transpose for periodic shuffling for each channel
  13. input_split = tf.split(inputs, channel_target, axis=3)
  14. output = tf.concat([phaseShift(x, scale, shape_1, shape_2) for x in input_split], axis=3)
  15. return output
  16. def phaseShift(inputs, scale, shape_1, shape_2):
  17. # Tackle the condition when the batch is None
  18. X = tf.reshape(inputs, shape_1)
  19. X = tf.transpose(X, [0, 1, 3, 2, 4])
  20. return tf.reshape(X, shape_2)

 10.生成器

本实验使用的基本模型是GAN,在GAN的生成器部分将从低分辨率的输入产生高分辨率的模型输出。

  1. class Generator(tf.keras.Model):
  2. def __init__(self, data_format='channels_last'):
  3. super(Generator, self).__init__(name='')
  4. if data_format == 'channels_first':
  5. self._input_shape = [-1, 3, 32, 32]
  6. self.bn_axis = 1
  7. else:
  8. assert data_format == 'channels_last'
  9. self._input_shape = [-1, 32, 32, 3]
  10. self.bn_axis = 3
  11. self.conv1 = tf.keras.layers.Conv2D(
  12. 64, kernel_size=9, strides=1, padding='SAME', data_format=data_format)
  13. self.prelu1 = tf.keras.layers.PReLU(shared_axes=[1, 2])
  14. self.res_blocks = [_IdentityBlock(64, 1, data_format) for _ in range(16)]
  15. self.conv2 = tf.keras.layers.Conv2D(
  16. 64, kernel_size=3, strides=1, padding='SAME', data_format=data_format)
  17. self.upconv1 = tf.keras.layers.Conv2D(
  18. 256, kernel_size=3, strides=1, padding='SAME', data_format=data_format)
  19. self.prelu2 = tf.keras.layers.PReLU(shared_axes=[1, 2])
  20. self.upconv2 = tf.keras.layers.Conv2D(
  21. 256, kernel_size=3, strides=1, padding='SAME', data_format=data_format)
  22. self.prelu3 = tf.keras.layers.PReLU(shared_axes=[1, 2])
  23. self.conv4 = tf.keras.layers.Conv2D(
  24. 3, kernel_size=9, strides=1, padding='SAME', data_format=data_format)
  25. def call(self, inputs):
  26. x = tf.reshape(inputs, self._input_shape)
  27. x = self.conv1(x)
  28. # x = tf.nn.leaky_relu(x)
  29. x = self.prelu1(x)
  30. x_start = x
  31. for i in range(len(self.res_blocks)):
  32. x = self.res_blocks[i](x)
  33. x = self.conv2(x)
  34. x = x + x_start
  35. x = self.upconv1(x)
  36. x = pixelShuffler(x)
  37. x = self.prelu2(x)
  38. x = self.upconv2(x)
  39. x = pixelShuffler(x)
  40. x = self.prelu3(x)
  41. x = self.conv4(x)
  42. x = tf.nn.tanh(x)
  43. return x

11.判别器

本实验的GAN判别器的输入是一张128*128的图像,目标输出是一个布尔值,也就是判断输入的图像是真的图像还是通过模型为伪造的图像。本实验设计的生成器是由全卷积网络实现。判别器的代码如下所示:

  1. class Discriminator(tf.keras.Model):
  2. def __init__(self, data_format='channels_last'):
  3. super(Discriminator, self).__init__(name='')
  4. if data_format == 'channels_first':
  5. self._input_shape = [-1, 3, 128, 128]
  6. self.bn_axis = 1
  7. else:
  8. assert data_format == 'channels_last'
  9. self._input_shape = [-1, 128, 128, 3]
  10. self.bn_axis = 3
  11. self.conv1 = tf.keras.layers.Conv2D(
  12. 64, kernel_size=3, strides=1, padding='SAME', data_format=data_format)
  13. self.conv2 = tf.keras.layers.Conv2D(
  14. 64, kernel_size=3, strides=2, padding='SAME', data_format=data_format)
  15. self.conv3 = tf.keras.layers.Conv2D(
  16. 128, kernel_size=3, strides=1, padding='SAME', data_format=data_format)
  17. # self.bn3 = tf.keras.layers.BatchNormalization(axis=self.bn_axis)
  18. self.conv4 = tf.keras.layers.Conv2D(
  19. 128, kernel_size=3, strides=2, padding='SAME', data_format=data_format)
  20. # self.bn4 = tf.keras.layers.BatchNormalization(axis=self.bn_axis)
  21. self.conv5 = tf.keras.layers.Conv2D(
  22. 256, kernel_size=3, strides=1, padding='SAME', data_format=data_format)
  23. # self.bn5 = tf.keras.layers.BatchNormalization(axis=self.bn_axis)
  24. self.conv6 = tf.keras.layers.Conv2D(
  25. 256, kernel_size=3, strides=2, padding='SAME', data_format=data_format)
  26. # self.bn6 = tf.keras.layers.BatchNormalization(axis=self.bn_axis)
  27. self.conv7 = tf.keras.layers.Conv2D(
  28. 512, kernel_size=3, strides=1, padding='SAME', data_format=data_format)
  29. # self.bn7 = tf.keras.layers.BatchNormalization(axis=self.bn_axis)
  30. self.conv8 = tf.keras.layers.Conv2D(
  31. 512, kernel_size=3, strides=2, padding='SAME', data_format=data_format)
  32. # self.bn8 = tf.keras.layers.BatchNormalization(axis=self.bn_axis)
  33. self.fc1 = tf.keras.layers.Dense(1024)
  34. self.fc2 = tf.keras.layers.Dense(1)
  35. def call(self, inputs):
  36. x = tf.reshape(inputs, self._input_shape)
  37. x = self.conv1(x)
  38. x = tf.nn.leaky_relu(x)
  39. x = self.conv2(x)
  40. # x = self.bn2(x)
  41. x = tf.nn.leaky_relu(x)
  42. x = self.conv3(x)
  43. # x = self.bn3(x)
  44. x = tf.nn.leaky_relu(x)
  45. x = self.conv4(x)
  46. # x = self.bn4(x)
  47. x = tf.nn.leaky_relu(x)
  48. x = self.conv5(x)
  49. # x = self.bn5(x)
  50. x = tf.nn.leaky_relu(x)
  51. x = self.conv6(x)
  52. # x = self.bn6(x)
  53. x = tf.nn.leaky_relu(x)
  54. x = self.conv7(x)
  55. # x = self.bn7(x)
  56. x = tf.nn.leaky_relu(x)
  57. x = self.conv8(x)
  58. # x = self.bn8(x)
  59. x = tf.nn.leaky_relu(x)
  60. # x = self.flatten(x)
  61. x = self.fc1(x)
  62. x = tf.nn.leaky_relu(x)
  63. x = self.fc2(x)
  64. # x = tf.nn.sigmoid(x)
  65. return x

12.损失函数与优化器定义

常规的GAN中,生成器的损失函数为对抗损失,定义为生成让判别器无法区分的数据分布,即让判别器将生成器生成的图像判定为真实图像的概率尽可能的高。但是在超分辨率任务中,这样的损失定义很难帮助生成器去生成细节足够真实的图像。因此本实验为生成器添加了额外的内容损失。内容损失的定义有两种方式,一种是经典的均方误差损失,即对生成器生成的网络与真实图像直接求均方误差,这样方式可以得到很高的信噪比,但是图像在高频细节上有缺失。第二种内容损失的基础是预训练的VGG 19网络的ReLU激活层为基础的VGG loss,然后通过求生成图像和原始图像特征表示的欧氏距离来计算当前的内容损失。

本实验选择了VGG loss作为内容损失,最终的生成器损失定义为内容损失和对抗损失的加权和。本部分该部分代码如下所示,注意首先定义了用于计算内容损失的VGG 19网络:

  1. def vgg19():
  2. vgg19 = tf.keras.applications.vgg19.VGG19(include_top=False, weights='imagenet', input_shape=(128, 128, 3))
  3. vgg19.trainable = False
  4. for l in vgg19.layers:
  5. l.trainable = False
  6. loss_model = tf.keras.Model(inputs=vgg19.input, outputs=vgg19.get_layer('block5_conv4').output)
  7. loss_model.trainable = False
  8. return loss_model

判别器损失与传统的GAN判别器损失类似,目标是将生成器生成的伪造图像尽可能判定为假的,将真实的原始图像尽可能判断为真的。最终的判别器损失是两部分的损失之和,如下所示:

  1. def create_g_loss(d_output, g_output, labels, loss_model):
  2. gene_ce_loss = tf.losses.sigmoid_cross_entropy(tf.ones_like(d_output), d_output)
  3. vgg_loss = tf.keras.backend.mean(tf.keras.backend.square(loss_model(labels) - loss_model(g_output)))
  4. g_loss = vgg_loss + 1e-3 * gene_ce_loss
  5. # g_loss = mse_loss + 1e-3 * gene_ce_loss
  6. return g_loss
  7. def create_d_loss(disc_real_output, disc_fake_output):
  8. disc_real_loss = tf.losses.sigmoid_cross_entropy(tf.ones_like(disc_real_output), disc_real_output)
  9. disc_fake_loss = tf.losses.sigmoid_cross_entropy(tf.zeros_like(disc_fake_output), disc_fake_output)
  10. disc_loss = tf.add(disc_real_loss, disc_fake_loss)
  11. return disc_loss
  12. #优化器选择的是Adam,beta 1设为0.9,beta 2设为0.999,epsilon设为1e-8,以减少震荡。代码如下所示:
  13. def create_optimizers():
  14. g_optimizer = tf.train.AdamOptimizer(learning_rate=1e-4, beta1=0.9, beta2=0.999, epsilon=1e-8)
  15. d_optimizer = tf.train.AdamOptimizer(learning_rate=1e-4, beta1=0.9, beta2=0.999, epsilon=1e-8)
  16. return g_optimizer, d_optimizer

在每一步的训练中,首先通过生成器获得伪造的高分辨率图像,然后分别计算生成器与判别器的损失,再分别更新生成器与判别器的参数,注意这里生成器与判别器的训练次数是1:1。代码如下所示:

  1. def train_step(feature, label, loss_model, generator, discriminator, g_optimizer, d_optimizer):
  2. with tf.GradientTape() as g_tape, tf.GradientTape() as d_tape:
  3. generated_images = generator(feature)
  4. real_output = discriminator(label)
  5. generated_output = discriminator(generated_images)
  6. g_loss = create_g_loss(generated_output, generated_images, label, loss_model)
  7. d_loss = create_d_loss(real_output, generated_output)
  8. gradients_of_generator = g_tape.gradient(g_loss, generator.variables)
  9. gradients_of_discriminator = d_tape.gradient(d_loss, discriminator.variables)
  10. g_optimizer.apply_gradients(zip(gradients_of_generator, generator.variables))
  11. d_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.variables))
  12. return g_loss, d_loss

13.训练

在上文中已经定义了实验的输入输出数据、生成器与判别器模型以及相应的优化器,本小节将介绍训练的过程。实验将首先导入上文的数据、模型与判别器,然后定义了checkpoint以持久化保存模型,然后一个批次一个批次地读取数据,进行每一步的训练。总体训练流程如下所示:

  1. def train_model(train_log_dir, train_image_ds, test_image_ds, epochs, checkpoint_dir):
  2. generator = Generator()
  3. discriminator = Discriminator()
  4. g_optimizer, d_optimizer = create_optimizers()
  5. checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
  6. checkpoint = tf.train.Checkpoint(g_optimizer=g_optimizer,
  7. d_optimizer=d_optimizer,
  8. generator=generator,
  9. discriminator=discriminator)
  10. loss_model = vgg19()
  11. for epoch in range(epochs):
  12. all_g_cost = all_d_cost = 0
  13. step = 0
  14. it = iter(train_image_ds)
  15. while True:
  16. try:
  17. image_batch, label_batch = next(it)
  18. step = step + 1
  19. g_loss, d_loss = train_step(image_batch, label_batch, loss_model, generator, discriminator,
  20. g_optimizer, d_optimizer)
  21. all_g_cost = all_g_cost + g_loss
  22. all_d_cost = all_d_cost + d_loss
  23. print("g_loss:",g_loss.numpy(), "d_loss:",d_loss.numpy())
  24. except StopIteration:
  25. break
  26. print("all_g_cost:",all_g_cost.numpy(), "all_d_cost:",all_d_cost.numpy())
  27. generate_and_save_images(train_log_dir, generator, epoch + 1, test_image_ds)
  28. # saving (checkpoint) the model every epoch
  29. checkpoint.save(file_prefix=checkpoint_prefix)
  30. def generate_and_save_images(train_dir, model, epoch, test_image_ds):
  31. dir = train_dir + str(epoch) + '/'
  32. feature_batch, _ = next(iter(test_image_ds))
  33. if tf.gfile.Exists(dir):
  34. tf.gfile.DeleteRecursively(dir)
  35. tf.gfile.MakeDirs(dir)
  36. predictions = model(feature_batch)
  37. for i, pred in enumerate(predictions):
  38. if i > 5:
  39. break
  40. imageio.imwrite(dir + 'image_{:02d}.png'.format(i), pred)
  1. # os.environ["CUDA_VISIBLE_DEVICES"] = '1'
  2. # config = tf.ConfigProto()
  3. # # config.gpu_options.per_process_gpu_memory_fraction = 0.8
  4. # config.gpu_options.allow_growth = True
  5. # tf.keras.backend.set_session(tf.Session(config=config))
  6. CHECKPOINT_DIR = 'checkpoint/'
  7. TRAIN_LOG_DIR = 'train_log/'
  8. train_hr_dir = 'data/'
  9. test_hr_dir = 'data-test/'
  10. resize_train_dir = 'train/'
  11. resize_test_dir = 'test/'
  12. tf.enable_eager_execution()
  13. def resize_data():
  14. # 因硬件配置低,只选择前100条图片训练
  15. train_files = tf.gfile.ListDirectory(train_hr_dir)[:100]
  16. test_files = tf.gfile.ListDirectory(test_hr_dir)[:100]
  17. delete_or_makedir(resize_train_dir)
  18. delete_or_makedir(resize_test_dir)
  19. for file in train_files:
  20. resize_image(file, train_hr_dir, resize_train_dir)
  21. for file in test_files:
  22. resize_image(file, test_hr_dir, resize_test_dir)
  23. print("resize done")
  24. def train():
  25. print("prepare train dirs")
  26. prepare_train_dirs(CHECKPOINT_DIR, TRAIN_LOG_DIR, True)
  27. print("load data")
  28. train_image_ds = load_data(resize_train_dir, training=True)
  29. test_image_ds = load_data(resize_test_dir, training=False)
  30. save_feature_label(TRAIN_LOG_DIR, test_image_ds)
  31. print("save feature label complete")
  32. # 这里训练次数设为1,可增加到60
  33. print("train begin")
  34. train_model(TRAIN_LOG_DIR, train_image_ds, test_image_ds, 1, CHECKPOINT_DIR)
  35. print("train done")
  1. #首次运行需要准备数据,后续不需要重复运行此代码
  2. resize_data()

  1. #开启训练
  2. train()

 

CPU训练的结果完全不行,完全看不清人脸。 

 14.实验评估(待实现)

本实验的训练数据有10656个图片,测试数据为5个图片。由于实验评估图像分辨率的提升效果,因此不需要按照一定比例划分训练集和测试集,而越多的训练数据也有更好的建模效果。但是需要注意的是,按照本实验现在的模型,需要的显存约为6.8G,如果增大批次数量、增加卷积层特征图数量、加深网络或增大原始图片分辨率,将进一步增加显存。

应用TensorFlow GPU版本,GPU是NVIDIA GeForce GTX 1080,显存大小为8G。一次迭代约需要7—8分钟,20次迭代的训练将持续约2.5小时。经过60次迭代的训练损失如图所示。

 

图中生成器的损失在前20次迭代训练下降,20次迭代到60次迭代后在0.02上下波动,损失进一步下降变得缓慢;判别器的损失下降较为明显,但是注意到有多次震荡情况的出现。本实验继续增大迭代次数可能会有更好的损失表示。

下图显示了迭代42次和迭代60次后的模型在测试集上的表现对比。

 

 从中可见模型对分辨率有非常明显的提升,第一列低分辨率图像在放大后细节部分非常模糊,而第二列迭代42次于第三列迭代60次后,图像在一些细节部分例如头发、五官等有了更加清晰的轮廓。但是模型与原始图相比仍然还有差距,不能做到以假乱真。一方面是因为实验只是在迭代60次后的模型上进行评估;一方面是因为原图是以JPG格式存储的,所以相比PNG无损格式的图片在细节上有先天的不足;还有原因是本实验将图片剪裁到了128*128,这是考虑到显存的权宜之计,更高的原始分辨率会有更丰富的细节信息。进一步增加迭代次数、进一步调整参数、扩大图片原始尺寸、使用原始PNG格式的图片等方法会导致最后的模型有更佳的效果。

总结:

  1. 使用CelebA数据集,在数据预处理后,构建以CNN、GAN和ResNet的混合模型,在迭代训练60次后有明显效果。

  2. 限于GPU算力、训练时间等因素,模型的输出相比原图不够完美

  3. 改进思路包括使用更高配GPU、调大图像原始尺寸、使用原始PNG格式图像、增加迭代次数、进一步调参等。

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/很楠不爱3/article/detail/316620?site
推荐阅读
相关标签
  

闽ICP备14008679号