autoencoder.py

#! /usr/bin/env python
# coding=utf-8
#================================================================
#   Copyright (C) 2019 * Ltd. All rights reserved.
#
#   Editor      : VIM
#   File name   : autoencoder.py
#   Author      : YunYang1994
#   Created date: 2019-11-05 19:53:58
#   Description :
#
#================================================================

import random
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from matplotlib import cm


class Encoder(tf.keras.Model):
    def __init__(self):
        super(Encoder, self).__init__()
        self.dense_1 = tf.keras.layers.Dense(128, activation='tanh')
        self.dense_2 = tf.keras.layers.Dense(64 , activation='tanh')
        self.dense_3 = tf.keras.layers.Dense(32 , activation='tanh')
        self.dense_4 = tf.keras.layers.Dense(2  , activation='sigmoid')
    def call(self, x, training=False):
        out = self.dense_1(x)
        out = self.dense_2(out)
        out = self.dense_3(out)
        out = self.dense_4(out)
        return out


class Decoder(tf.keras.Model):
    def __init__(self):
        super(Decoder, self).__init__()
        self.dense_1 = tf.keras.layers.Dense(16 , activation='tanh')
        self.dense_2 = tf.keras.layers.Dense(64 , activation='tanh')
        self.dense_3 = tf.keras.layers.Dense(128, activation='tanh')
        self.dense_4 = tf.keras.layers.Dense(784, activation='sigmoid')
    def call(self, x, training=False):
        out = self.dense_1(x)
        out = self.dense_2(out)
        out = self.dense_3(out)
        out = self.dense_4(out)
        return out

class Autoencoder(tf.keras.Model):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()
    def call(self, x, training=False):
        out = self.encoder(x, training)
        out = self.decoder(out, training)
        return out


(X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()
X_train = np.reshape(X_train, (-1, 784)) / 255.
batch_size = 512
optimizer = tf.keras.optimizers.Adam(lr=0.001)
model = Autoencoder()
sample = np.reshape(X_test[:5], (5, 784))

for step in range(10000):
    true_image = X_train[np.random.choice(X_train.shape[0], batch_size)]
    with tf.GradientTape() as tape:
        pred_image = model(true_image, training=True)
        loss = tf.reduce_mean(tf.square(pred_image-true_image))
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    if step % 200  == 0: print("=> %4d loss %.4f" %(step, loss))
    if step % 1000 == 0:
        pred_image = model(sample, training=False)
        pred_image = np.reshape(pred_image, (5, 28, 28))
        show_image = np.concatenate(pred_image[:5], -1)
        plt.tight_layout()
        plt.imshow(show_image)
        plt.savefig("%d.png" %((step+1) // 2000))

"""
visualize embedding in 2D
"""
sample = np.reshape(X_test[:5000], (5000, 784))
label = y_test[:5000]
embeddings = model.encoder(sample, training=False)

fig,ax = plt.subplots()
X, Y = embeddings[:,0].numpy(), embeddings[:,1].numpy()
ax.set_xlim(X.min(), X.max())
ax.set_ylim(Y.min(), Y.max())
for x,y,l in zip(X,Y,label):
    c = cm.rainbow(int(255 *l/ 9))
    ax.text(x, y, l, color=c)
    # plt.plot(x,y, '.', c=c)
plt.axis('off')
plt.legend()
plt.tight_layout()
plt.savefig("embedding.png")
plt.show()