Tensorflow2.0实现自编码器AE

关于自编码器的结构，我都不介绍了，比较简单，框架结构如下：

 代码：

import tensorflow as tf
import tensorflow.keras.datasets.mnist as minst
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from tensorflow.keras import layers
from tensorflow.keras.models import Model
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')
tf.test.is_gpu_available()

(train_image, train_label_or), (test_image, test_label_or) = minst.load_data()

#对手写数字的标签进行one_hat编码
train_label = np.eye(10)[train_label_or]

train_image = tf.cast(train_image,tf.float32)

#将数据进行归一化操作
train_image = train_image/255.0

images_count = train_image.shape[0]

#(60000, 28, 28) ----> (60000, 28, 28, 1)
train_image = np.expand_dims(train_image, -1)
test_image = np.expand_dims(test_image, -1)

BATCH_SIZE = 256

ts_train_images = tf.data.Dataset.from_tensor_slices(train_image)
ts_train_labels = tf.data.Dataset.from_tensor_slices(train_label)
ts_train_set = tf.data.Dataset.zip((ts_train_images,ts_train_labels))

train_dataset = ts_train_set.shuffle(images_count).batch(BATCH_SIZE)

def encoder():
    encoder_input = layers.Input(shape = (28,28,1))
    
    x = layers.Flatten()(encoder_input)
    
    x = layers.Dense(1024)(x)
    x = layers.BatchNormalization()(x)
    x = layers.LeakyReLU()(x)
    
    x = layers.Dense(512)(x)
    x = layers.BatchNormalization()(x)
    x = layers.LeakyReLU()(x)
    
    encoder_output = layers.Dense(10)(x)
    encoder_model = Model(inputs = encoder_input,outputs = encoder_output)
    print(encoder_model.summary())
    return encoder_model
Encoder = encoder()

def decoder():  
    decoder_input = layers.Input(shape = (10,))
    decoder_input1 = layers.Dense(7*7*32)(decoder_input)
    decoder_input2 = layers.Activation("relu")(decoder_input1)
    decoder_input3 = layers.Reshape(target_shape=(7, 7, 32))(decoder_input2)
    
    y = layers.Conv2DTranspose(
              filters=64,
              kernel_size=3,
              strides=(2, 2),
              padding="SAME",
              )(decoder_input3)
    y = layers.Dropout(rate = 0.7)(y)
    y = layers.Activation("relu")(y)
    
    y = layers.Conv2DTranspose(
              filters=32,
              kernel_size=3,
              strides=(2, 2),
              padding="SAME",
              )(y)
    y = layers.Dropout(rate = 0.7)(y)
    y = layers.Activation("relu")(y)
    
    y = layers.Conv2DTranspose(
              filters=1, kernel_size=3, strides=(1, 1), padding="SAME")(y)             
    y = layers.BatchNormalization()(y)
    decoder_output = layers.Activation("sigmoid")(y)
    decoder_model = Model(inputs = decoder_input,outputs = decoder_output)
    print(decoder_model.summary())
    return decoder_model
Decoder = decoder()


optimizer = tf.keras.optimizers.Adam(1e-3)

def train_epoch(image_batch,label_batch):
    with tf.GradientTape() as encoder_tap,tf.GradientTape() as decoder_tap:
        img_fea = Encoder(image_batch,training=True)
        res_img = Decoder(img_fea,training=True)
        
        sub_per = res_img - image_batch
        restructure_loss = tf.reduce_mean(tf.norm(tf.reshape(sub_per,shape=(sub_per.shape[0],-1)),2,axis=1))
        
        
        encoder_gard = encoder_tap.gradient(restructure_loss,Encoder.trainable_variables)
        decoder_grad = decoder_tap.gradient(restructure_loss,Decoder.trainable_variables)
        
        optimizer.apply_gradients(zip(encoder_gard,Encoder.trainable_variables))    
        optimizer.apply_gradients(zip(decoder_grad,Decoder.trainable_variables))
        
        return restructure_loss.numpy()

def generator():
    input_img = train_image[np.random.randint(train_image.shape[0])]
    input_img = np.expand_dims(input_img,0)
    res_img = Encoder.predict(input_img)
    img = Decoder.predict(res_img)
    img = tf.sigmoid(img)
    plt.figure(figsize = (2,3))
    ax1 = plt.subplot(1,2,1)
    ax1.axis("off")
    ax1.imshow(input_img[0],cmap = "gray")
    ax2 = plt.subplot(1,2,2)
    ax2.axis("off")
    ax2.imshow(img[0],cmap = "gray")
    plt.show()


"""
迭代的次数
"""
EPOCHS = 100

if __name__ == "__main__":
    train() 

 可以看出来，训练几个epochs，重构的效果就比较好了。