Python
class DownConvLayer(tf.keras.layers.Layer):
def __init__(self, dim):
super(DownConvLayer, self).__init__()
self.conv = tf.keras.layers.Conv2D(dim, 3, activation=tf.keras.layers.ReLU(), use_bias=False, padding='same')
self.pool = tf.keras.layers.MaxPool2D(2)
def call(self, x, training=False, **kwargs):
x = self.conv(x)
x = self.pool(x)
return x
class UpConvLayer(tf.keras.layers.Layer):
def __init__(self, dim):
super(UpConvLayer, self).__init__()
self.conv = tf.keras.layers.Conv2D(dim, 3, activation=tf.keras.layers.ReLU(), use_bias=False, padding='same')
# 通过UpSampling2D上采样
self.pool = tf.keras.layers.UpSampling2D(2)
def call(self, x, training=False, **kwargs):
x = self.conv(x)
x = self.pool(x)
return x
# 示例代码都是通过非常简单的卷积操作实现编码器和解码器
class Encoder(tf.keras.layers.Layer):
def __init__(self, dim, layer_num=3):
super(Encoder, self).__init__()
self.convs = [DownConvLayer(dim) for _ in range(layer_num)]
def call(self, x, training=False, **kwargs):
for conv in self.convs:
x = conv(x, training)
return x
class Decoder(tf.keras.layers.Layer):
def __init__(self, dim, layer_num=3):
super(Decoder, self).__init__()
self.convs = [UpConvLayer(dim) for _ in range(layer_num)]
self.final_conv = tf.keras.layers.Conv2D(1, 3, strides=1)
def call(self, x, training=False, **kwargs):
for conv in self.convs:
x = conv(x, training)
# 将图像转成和输入图像shape一致
reconstruct = self.final_conv(x)
return reconstruct
class AutoEncoderModel(tf.keras.Model):
def __init__(self):
super(AutoEncoderModel, self).__init__()
self.encoder = Encoder(64, layer_num=3)
self.decoder = Decoder(64, layer_num=3)
def call(self, inputs, training=None, mask=None):
image = inputs[0]
# 得到图像的特征表示
latent = self.encoder(image, training)
# 通过特征重建图像
reconstruct_img = self.decoder(latent, training)
return reconstruct_img
@tf.function
def train_step(self, data):
img = data["image"]
with tf.GradientTape() as tape:
reconstruct_img = self((img,), True)
trainable_vars = self.trainable_variables
# 利用l2 loss 来判断重建图片和原始图像的一致性
l2_loss = (reconstruct_img - img) ** 2
l2_loss = tf.reduce_mean(tf.reduce_sum(
l2_loss, axis=(1, 2, 3)
))
gradients = tape.gradient(l2_loss, trainable_vars)
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
return {"l2_loss": l2_loss} |