#
#作者:韦访
#博客:https://blog.csdn.net/rookie_wei
#微信:1007895847
#添加微信的备注一下是CSDN的
#欢迎大家一起学习
#
1、概述
上两讲中,我们了解了怎么将图像超分辨率模型SRGAN移植到安卓APP中,但是并没怎么涉及到SRGAN模型本身的知识,这一讲就来补补。源码怎么下载和使用,请看第22讲:
https://blog.csdn.net/rookie_wei/article/details/110356225
环境配置:
操作系统:Win10 64位
显卡:GTX 1080ti
Python:Python3.7
TensorFlow:2.3.0
2、DIV2K数据集
我们这里使用的是DIV2K数据集,源码中的data.py提供了下载和解压数据集的功能,我们先来调用它把数据集下载并解压到本地,再来看看这个数据集是什么回事,代码如下,
from data import DIV2K
train_loader = DIV2K(scale=4, # 2, 3, 4 or 8
downgrade='bicubic', # 'bicubic', 'unknown', 'mild' or 'difficult'
subset='train') # Training dataset are images 001 - 800
# Create a tf.data.Dataset
train_ds = train_loader.dataset(batch_size=16, # batch size as described in the EDSR and WDSR papers
random_transform=True, # random crop, flip, rotate as described in the EDSR paper
repeat_count=None) # repeat iterating over training images indefinitely
valid_loader = DIV2K(scale=4, # 2, 3, 4 or 8
downgrade='bicubic', # 'bicubic', 'unknown', 'mild' or 'difficult'
subset='valid') # Validation dataset are images 801 - 900
# Create a tf.data.Dataset
valid_ds = valid_loader.dataset(batch_size=1, # use batch size of 1 as DIV2K images have different size
random_transform=False, # use DIV2K images in original size
repeat_count=1) # 1 epoch
运行结果,
如果一切顺利的话,会在当前项目的根目录下生成一个名为“.div2k”的文件夹,该文件夹下有caches和images两个文件,我们先来看一下images文件夹,
如上图所示,images文件夹下一共有4个子文件夹,通过文件夹名字可以猜到,带“train”的就是我们的训练集,带“valid”的就是我们的验证集。先来对比DIV2K_train_HR和DIV2K_train_LR_bicubic这两个文件夹,DIV2K_train_HR文件夹下直接就是800张图片,而且图片都是高清的,
DIV2K_train_LR_bicubic文件夹下则有一个名为“X4”的子文件夹,子文件夹下也有800张图片,不过并不是高清图,
对比DIV2K_train_HR和DIV2K_train_LR_bicubic这两个文件夹可以发现,他们的图片其实是一样的,只不过DIV2K_train_LR_bicubic文件夹下的图片的分辨率是DIV2K_train_HR文件夹下对应的图片的1/4,而且是经过bicubic算法缩放的。
这样,我们就可以让DIV2K_train_LR_bicubic下的图片经过SRGAN模型还原成本身4倍大小的高清图,而对应的DIV2K_train_HR文件夹下的图片就是我们的“label”。
3、数据增强
上面我们通过DIV2K这个类将数据集下载并解压好了,现在,我们来看DIV2K这个类的代码,看看它对数据集进行怎样的预处理。
首先看我们怎么实例化这个类时传入的参数,
train_loader = DIV2K(scale=4, # 2, 3, 4 or 8
downgrade='bicubic', # 'bicubic', 'unknown', 'mild' or 'difficult'
subset='train') # Training dataset are images 001 - 800
上面代码中,scale指的是低分辨率图相对于高清图的缩小比例,我们这里用的是4倍。downgrade指的是缩放算法,我们用常见的bicubic算法。suset指定是训练集还是测试集。DIV2K的__init__函数就是初始化一些变量,比较简单就不看了,重点来看它的dataset函数,代码如下,
def dataset(self, batch_size=16, repeat_count=None, random_transform=True):
ds = tf.data.Dataset.zip((self.lr_dataset(), self.hr_dataset()))
if random_transform:
ds = ds.map(lambda lr, hr: random_crop(lr, hr, scale=self.scale), num_parallel_calls=AUTOTUNE)
ds = ds.map(random_rotate, num_parallel_calls=AUTOTUNE)
ds = ds.map(random_flip, num_parallel_calls=AUTOTUNE)
ds = ds.batch(batch_size)
ds = ds.repeat(repeat_count)
ds = ds.prefetch(buffer_size=AUTOTUNE)
return ds
可以看到,这个函数就是用tf.data.Dataset接口对数据进行数据增强和设置batch等一些操作,先来看lr_dataset函数,代码如下,
def lr_dataset(self):
if not os.path.exists(self._lr_images_dir()):
download_archive(self._lr_images_archive(), self.images_dir, extract=True)
ds = self._images_dataset(self._lr_image_files()).cache(self._lr_cache_file())
if not os.path.exists(self._lr_cache_index()):
self._populate_cache(ds, self._lr_cache_file())
return ds
如上代码所示,它先去检查数据集的目录是否存在,如果不存在的话就下载并解压数据集。接着,看_images_dataset函数,代码如下,
@staticmethod
def _images_dataset(image_files):
ds = tf.data.Dataset.from_tensor_slices(image_files)
ds = ds.map(tf.io.read_file)
ds = ds.map(lambda x: tf.image.decode_png(x, channels=3), num_parallel_calls=AUTOTUNE)
return ds
这个函数就是使用tf.data.Dataset接口创建Dataset对象的,然后,使用tf.io.read_file和tf.image.decode_png来读取和解析图片数据。image_files就是所有低分辨率图片的文件路径列表,通过_lr_image_files函数获取,代码如下,
def _lr_image_files(self):
images_dir = self._lr_images_dir()
return [os.path.join(images_dir, self._lr_image_file(image_id)) for image_id in self.image_ids]
再看一下lr_dataset函数,再得到Dataset对象以后,它又调用了Dataset的cache函数将数据缓存到cache文件中,也就是.div2k/cache文件夹下的文件,如下图示,
这样做估计能提高读取数据集的速度吧。
hr_dataset函数的操作和lr_dataset函数类似。
再回到dataset函数,得到lr和hr的Dataset对象以后,就要对它们进行数据增强操作了,先来看random_crop操作,代码如下,
def random_crop(lr_img, hr_img, hr_crop_size=96, scale=2):
lr_crop_size = hr_crop_size // scale
lr_img_shape = tf.shape(lr_img)[:2]
lr_w = tf.random.uniform(shape=(), maxval=lr_img_shape[1] - lr_crop_size + 1, dtype=tf.int32)
lr_h = tf.random.uniform(shape=(), maxval=lr_img_shape[0] - lr_crop_size + 1, dtype=tf.int32)
hr_w = lr_w * scale
hr_h = lr_h * scale
lr_img_cropped = lr_img[lr_h:lr_h + lr_crop_size, lr_w:lr_w + lr_crop_size]
hr_img_cropped = hr_img[hr_h:hr_h + hr_crop_size, hr_w:hr_w + hr_crop_size]
return lr_img_cropped, hr_img_cropped
这个函数的作用是随机裁剪低分辨率图和高清图“相同位置”的一小块图片。为什么不用整张图做训练?
- 在上一讲中我们知道了,运行SRGAN模型是非常耗内存的,如果用大图来训练机器硬件资源消耗过大。
- 随机截取图片的小片段来训练其实也是扩充了我们的数据集“图片”数量,防止过拟合。本来我们训练集里也才800张图片,现在这样随机截取图片,相当于图片扩充了N倍了。
接着回到dataset函数,随机裁剪操作以后,还有随机旋转和随机翻转的操作,这些都是数据增强中的常规操作了,就不多说了。
4、生成器pre-training
SRGAN是生成对抗神经网络(GAN),所以它由生成器和判别器组成。在进行GAN训练之前,一般可以先单独对生成器进行pre-training,得到一个“还不错”的生成器模型作为初始模型以后,再进行GAN训练,继续优化生成器和判别器。不只是GAN,有很多其他的模型也都会才有这种预训练的方式。先来看pre-training的代码,代码如下,
from model.srgan import generator
from train import SrganGeneratorTrainer
from data import DIV2K
train_loader = DIV2K(scale=4, # 2, 3, 4 or 8
downgrade='bicubic', # 'bicubic', 'unknown', 'mild' or 'difficult'
subset='train') # Training dataset are images 001 - 800
# Create a tf.data.Dataset
train_ds = train_loader.dataset(batch_size=16, # batch size as described in the EDSR and WDSR papers
random_transform=True, # random crop, flip, rotate as described in the EDSR paper
repeat_count=None) # repeat iterating over training images indefinitely
valid_loader = DIV2K(scale=4, # 2, 3, 4 or 8
downgrade='bicubic', # 'bicubic', 'unknown', 'mild' or 'difficult'
subset='valid') # Validation dataset are images 801 - 900
# Create a tf.data.Dataset
valid_ds = valid_loader.dataset(batch_size=1, # use batch size of 1 as DIV2K images have different size
random_transform=False, # use DIV2K images in original size
repeat_count=1) # 1 epoch
# Create a training context for the generator (SRResNet) alone.
pre_trainer = SrganGeneratorTrainer(model=generator(), checkpoint_dir=f'.ckpt/pre_generator')
# Pre-train the generator with 1,000,000 steps (100,000 works fine too).
pre_trainer.train(train_ds, valid_ds.take(10), steps=1000000, evaluate_every=1000)
# Save weights of pre-trained generator (needed for fine-tuning with GAN).
pre_trainer.model.save_weights('weights/srgan/pre_generator.h5')
5、生成器网络
上面代码中,generator()就是SRGAN的生成器网络,看代码之前,我们先来看原始论文中怎么设计这个网络的,
然后再看代码,
def upsample(x_in, num_filters):
x = Conv2D(num_filters, kernel_size=3, padding='same')(x_in)
x = Lambda(pixel_shuffle(scale=2))(x)
return PReLU(shared_axes=[1, 2])(x)
def res_block(x_in, num_filters, momentum=0.8):
x = Conv2D(num_filters, kernel_size=3, padding='same')(x_in)
x = BatchNormalization(momentum=momentum)(x)
x = PReLU(shared_axes=[1, 2])(x)
x = Conv2D(num_filters, kernel_size=3, padding='same')(x)
x = BatchNormalization(momentum=momentum)(x)
x = Add()([x_in, x])
return x
def sr_resnet(num_filters=64, num_res_blocks=16):
x_in = Input(shape=(None, None, 3))
x = Lambda(normalize_01)(x_in)
x = Conv2D(num_filters, kernel_size=9, padding='same')(x)
x = x_1 = PReLU(shared_axes=[1, 2])(x)
for _ in range(num_res_blocks):
x = res_block(x, num_filters)
x = Conv2D(num_filters, kernel_size=3, padding='same')(x)
x = BatchNormalization()(x)
x = Add()([x_1, x])
x = upsample(x, num_filters * 4)
x = upsample(x, num_filters * 4)
x = Conv2D(3, kernel_size=9, padding='same', activation='tanh')(x)
x = Lambda(denormalize_m11)(x)
return Model(x_in, x)
generator = sr_resnet
模型对输入图片的大小是不做限制的,论文中的模型结构已经很清楚了,照着写代码就可以了。
回到SrganGeneratorTrainer类,该类继承了Trainer类,具体的训练代码是由Trainer类实现的,之所以这样设计,是因为这个代码不只是实现了SRGAN,还实现了EDSR和WDSR模型,这两个模型也都是图像超分辨率的模型。先看SrganGeneratorTrainer类的实现,代码如下,
class SrganGeneratorTrainer(Trainer):
def __init__(self,
model,
checkpoint_dir,
learning_rate=1e-4):
super().__init__(model, loss=MeanSquaredError(), learning_rate=learning_rate, checkpoint_dir=checkpoint_dir)
def train(self, train_dataset, valid_dataset, steps=1000000, evaluate_every=1000, save_best_only=True):
super().train(train_dataset, valid_dataset, steps, evaluate_every, save_best_only)
这里主要是设置了一些超参数,重点来看Trainer类的train函数,代码如下,
def train(self, train_dataset, valid_dataset, steps, evaluate_every=1000, save_best_only=False):
loss_mean = Mean()
ckpt_mgr = self.checkpoint_manager
ckpt = self.checkpoint
self.now = time.perf_counter()
for lr, hr in train_dataset.take(steps - ckpt.step.numpy()):
ckpt.step.assign_add(1)
step = ckpt.step.numpy()
loss = self.train_step(lr, hr)
loss_mean(loss)
if step % evaluate_every == 0:
loss_value = loss_mean.result()
loss_mean.reset_states()
# Compute PSNR on validation dataset
psnr_value = self.evaluate(valid_dataset)
duration = time.perf_counter() - self.now
print(f'{step}/{steps}: loss = {loss_value.numpy():.3f}, PSNR = {psnr_value.numpy():3f} ({duration:.2f}s)')
if save_best_only and psnr_value <= ckpt.psnr:
self.now = time.perf_counter()
# skip saving checkpoint, no PSNR improvement
continue
ckpt.psnr = psnr_value
ckpt_mgr.save()
self.now = time.perf_counter()
ckpt_mgr和ckpt是用来在训练过程中保存模型和中断训练后重新加载模型用的,在__init__函数中定义,代码如下,
self.checkpoint = tf.train.Checkpoint(step=tf.Variable(0),
psnr=tf.Variable(-1.0),
optimizer=Adam(learning_rate),
model=model)
self.checkpoint_manager = tf.train.CheckpointManager(checkpoint=self.checkpoint,
directory=checkpoint_dir,
max_to_keep=3)
train_step函数是重点,代码如下,
@tf.function
def train_step(self, lr, hr):
with tf.GradientTape() as tape:
lr = tf.cast(lr, tf.float32)
hr = tf.cast(hr, tf.float32)
sr = self.checkpoint.model(lr, training=True)
loss_value = self.loss(hr, sr)
gradients = tape.gradient(loss_value, self.checkpoint.model.trainable_variables)
self.checkpoint.optimizer.apply_gradients(zip(gradients, self.checkpoint.model.trainable_variables))
return loss_value
首先将lr和hr的数据类型转成float型,然后,将lr的数据送入模型中,得到模型生成的高清图sr,然后经过loss函数求hr和sr的误差,通过SrganGeneratorTrainer类的__init__函数可以看到,这里的loss函数其实就是均方误差函数MeanSquaredError。得到损失值以后,再求梯度并使用优化器优化模型即可。
train完以后,将预训练的模型保存起来,作为GAN训练的生成器的初始值。
pre_trainer.model.save_weights('weights/srgan/pre_generator.h5')
6、判别器网络
上面讲了生成器网络,现在来看GAN中另一个重要的网络----判别器。先来看原论文中的网络结构,如下图所示,
代码如下,
def discriminator_block(x_in, num_filters, strides=1, batchnorm=True, momentum=0.8):
x = Conv2D(num_filters, kernel_size=3, strides=strides, padding='same')(x_in)
if batchnorm:
x = BatchNormalization(momentum=momentum)(x)
return LeakyReLU(alpha=0.2)(x)
def discriminator(num_filters=64):
x_in = Input(shape=(HR_SIZE, HR_SIZE, 3))
x = Lambda(normalize_m11)(x_in)
x = discriminator_block(x, num_filters, batchnorm=False)
x = discriminator_block(x, num_filters, strides=2)
x = discriminator_block(x, num_filters * 2)
x = discriminator_block(x, num_filters * 2, strides=2)
x = discriminator_block(x, num_filters * 4)
x = discriminator_block(x, num_filters * 4, strides=2)
x = discriminator_block(x, num_filters * 8)
x = discriminator_block(x, num_filters * 8, strides=2)
x = Flatten()(x)
x = Dense(1024)(x)
x = LeakyReLU(alpha=0.2)(x)
x = Dense(1, activation='sigmoid')(x)
return Model(x_in, x)
也是根据网络结构搭建就好了。
7、损失函数
在继续往下看代码之前,先讲我们待会需要用到的几个损失函数。
8、训练GAN
接下来,就是训练GAN 了,先看代码,
from model.srgan import generator, discriminator
from train import SrganTrainer
from data import DIV2K
train_loader = DIV2K(scale=4, # 2, 3, 4 or 8
downgrade='bicubic', # 'bicubic', 'unknown', 'mild' or 'difficult'
subset='train') # Training dataset are images 001 - 800
# Create a tf.data.Dataset
train_ds = train_loader.dataset(batch_size=16, # batch size as described in the EDSR and WDSR papers
random_transform=True, # random crop, flip, rotate as described in the EDSR paper
repeat_count=None) # repeat iterating over training images indefinitely
valid_loader = DIV2K(scale=4, # 2, 3, 4 or 8
downgrade='bicubic', # 'bicubic', 'unknown', 'mild' or 'difficult'
subset='valid') # Validation dataset are images 801 - 900
# Create a tf.data.Dataset
valid_ds = valid_loader.dataset(batch_size=1, # use batch size of 1 as DIV2K images have different size
random_transform=False, # use DIV2K images in original size
repeat_count=1) # 1 epoch
# Create a new generator and init it with pre-trained weights.
gan_generator = generator()
gan_generator.load_weights('weights/srgan/pre_generator.h5')
# Create a training context for the GAN (generator + discriminator).
gan_trainer = SrganTrainer(generator=gan_generator, discriminator=discriminator())
# Train the GAN with 200,000 steps.
gan_trainer.train(train_ds, steps=200000)
# Save weights of generator and discriminator.
gan_trainer.generator.save_weights('weights/srgan/gan_generator.h5')
gan_trainer.discriminator.save_weights('weights/srgan/gan_discriminator.h5')
首先,生成器先加载之前预训练过的模型,
# Create a new generator and init it with pre-trained weights.
gan_generator = generator()
gan_generator.load_weights('weights/srgan/pre_generator.h5')
然后实例化SrganTrainer类,该类的两个参数分别是加载预训练模型后的生成器和判别器,
# Create a training context for the GAN (generator + discriminator).
gan_trainer = SrganTrainer(generator=gan_generator, discriminator=discriminator())
重点来看SrganTrainer类的train函数,
def train(self, train_dataset, steps=200000):
pls_metric = Mean()
dls_metric = Mean()
step = 0
for lr, hr in train_dataset.take(steps):
step += 1
pl, dl = self.train_step(lr, hr)
pls_metric(pl)
dls_metric(dl)
if step % 50 == 0:
print(f'{step}/{steps}, perceptual loss = {pls_metric.result():.4f}, discriminator loss = {dls_metric.result():.4f}')
pls_metric.reset_states()
dls_metric.reset_states()
重点来看train_step函数,
@tf.function
def train_step(self, lr, hr):
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
lr = tf.cast(lr, tf.float32)
hr = tf.cast(hr, tf.float32)
sr = self.generator(lr, training=True)
hr_output = self.discriminator(hr, training=True)
sr_output = self.discriminator(sr, training=True)
con_loss = self._content_loss(hr, sr)
gen_loss = self._generator_loss(sr_output)
perc_loss = con_loss + 0.001 * gen_loss
disc_loss = self._discriminator_loss(hr_output, sr_output)
gradients_of_generator = gen_tape.gradient(perc_loss, self.generator.trainable_variables)
gradients_of_discriminator = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)
self.generator_optimizer.apply_gradients(zip(gradients_of_generator, self.generator.trainable_variables))
self.discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, self.discriminator.trainable_variables))
return perc_loss, disc_loss
@tf.function
def _content_loss(self, hr, sr):
sr = preprocess_input(sr)
hr = preprocess_input(hr)
sr_features = self.vgg(sr) / 12.75
hr_features = self.vgg(hr) / 12.75
return self.mean_squared_error(hr_features, sr_features)
这里其实就是先求hr和sr的VGG特征图,然后再求特征图之间的MSE,这里之所以要除以12.75,是因为论文里说的,
继续回到train_step函数,计算内容损失以后,再通过_generator_loss函数计算对抗损失,代码如下,
def _generator_loss(self, sr_out):
return self.binary_cross_entropy(tf.ones_like(sr_out), sr_out)
这里就是“欺骗”判别结果的代码,sr_out都是重建的图像,我们这里却说它都是原始高清图像,如果判别器没有被“骗”到了,那就算到生成器的损失里。继续看train_step函数,
perc_loss = con_loss + 0.001 * gen_loss
就是我们要求的感知损失,
disc_loss = self._discriminator_loss(hr_output, sr_output)
则是求判别器的损失,我们来看_discriminator_loss函数,
def _discriminator_loss(self, hr_out, sr_out):
hr_loss = self.binary_cross_entropy(tf.ones_like(hr_out), hr_out)
sr_loss = self.binary_cross_entropy(tf.zeros_like(sr_out), sr_out)
return hr_loss + sr_loss
因为这是求判别器的损失,所以这里就得老实的告诉判别器,到底哪些图是“真图”,哪些图是“赝品”,这样判别器才能更好的工作。接下来的代码就是求生成器和判别器的梯度和优化了。
gradients_of_generator = gen_tape.gradient(perc_loss, self.generator.trainable_variables)
gradients_of_discriminator = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)
self.generator_optimizer.apply_gradients(zip(gradients_of_generator, self.generator.trainable_variables))
self.discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, self.discriminator.trainable_variables))