Autoencoders TF

# Bundled, deterministic handwritten-digits fixture.
from sklearn.datasets import load_digits as _course_load_digits

def _course_digits_load_data():
    course_digits = _course_load_digits()
    images = np.kron(
        course_digits.images.astype(np.float32),
        np.ones((3, 3), dtype=np.float32),
    )
    images = np.pad(images, ((0, 0), (2, 2), (2, 2))) * (255.0 / 16.0)
    labels = np.asarray(course_digits.target, dtype=np.int64)
    split = int(0.8 * len(images))
    return (images[:split], labels[:split]), (images[split:], labels[split:])

import tensorflow as tf
from tensorflow.keras.datasets import mnist
import numpy as np
import matplotlib.pyplot as plt

(x_train, y_trainclass), (x_test, y_testclass) = _course_digits_load_data()

def plotn(n,x):
  fig,ax = plt.subplots(1,n)
  for i,z in enumerate(x[0:n]):
    ax[i].imshow(z.reshape(28,28) if z.size==28*28 else z.reshape(14,14) if z.size==14*14 else z)
  plt.show()

plotn(5,x_train)

from tensorflow.keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D, Lambda
from tensorflow.keras.models import Model
from tensorflow.keras.losses import binary_crossentropy,mse

input_img = Input(shape=(28, 28, 1))

x = Conv2D(16, (3, 3), activation='relu', padding='same')(input_img)
x = MaxPooling2D((2, 2), padding='same')(x)
x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
x = MaxPooling2D((2, 2), padding='same')(x)
x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
encoded = MaxPooling2D((2, 2), padding='same')(x)

encoder = Model(input_img,encoded)

input_rep = Input(shape=(4,4,8))

x = Conv2D(8, (3, 3), activation='relu', padding='same')(input_rep)
x = UpSampling2D((2, 2))(x)
x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
x = UpSampling2D((2, 2))(x)
x = Conv2D(16, (3, 3), activation='relu')(x)
x = UpSampling2D((2, 2))(x)
decoded = Conv2D(1, (3, 3), activation='sigmoid', padding='same')(x)

decoder = Model(input_rep,decoded)

autoencoder = Model(input_img, decoder(encoder(input_img)))
autoencoder.compile(optimizer='adam', loss='binary_crossentropy')

x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.
x_train = np.reshape(x_train, (len(x_train), 28, 28, 1))
x_test = np.reshape(x_test, (len(x_test), 28, 28, 1))

autoencoder.fit(x_train, x_train,
                epochs=2,
                batch_size=128,
                shuffle=True,
                validation_data=(x_test, x_test))

y_test = autoencoder.predict(x_test[0:5])
plotn(5,x_test)
plotn(5,y_test)

encoder = Model(input_img, encoded)
encoded_imgs = encoder.predict(x_test[0:5])

plotn(5,encoded_imgs.reshape(5,-1,8))

print(encoded_imgs.max(),encoded_imgs.min())
res = decoder.predict(7*np.random.rand(7,4,4,8))
plotn(7,res)

def noisify(data):
  return np.clip(data+np.random.normal(loc=0.5,scale=0.5,size=data.shape),0.,1.)

x_train_noise = noisify(x_train)
x_test_noise = noisify(x_test)

plotn(5,x_train_noise)

autoencoder.fit(x_train_noise, x_train,
                epochs=2,
                batch_size=128,
                shuffle=True,
                validation_data=(x_test_noise, x_test))

y_test = autoencoder.predict(x_test_noise[0:5])
plotn(5,x_test_noise)
plotn(5,y_test)

x_train_lr = tf.keras.layers.AveragePooling2D(pool_size=(2, 2), strides=(2, 2))(x_train).numpy()
x_test_lr = tf.keras.layers.AveragePooling2D(pool_size=(2, 2), strides=(2, 2))(x_test).numpy()
plotn(5,x_train_lr)

from tensorflow.keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D, Lambda
from tensorflow.keras.models import Model
from tensorflow.keras.losses import binary_crossentropy,mse

input_img = Input(shape=(14, 14, 1))

x = Conv2D(16, (3, 3), activation='relu', padding='same')(input_img)
x = MaxPooling2D((2, 2), padding='same')(x)
x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
encoded = MaxPooling2D((2, 2), padding='same')(x)

encoder = Model(input_img,encoded)

input_rep = Input(shape=(4,4,8))

x = Conv2D(8, (3, 3), activation='relu', padding='same')(input_rep)
x = UpSampling2D((2, 2))(x)
x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
x = UpSampling2D((2, 2))(x)
x = Conv2D(16, (3, 3), activation='relu')(x)
x = UpSampling2D((2, 2))(x)
decoded = Conv2D(1, (3, 3), activation='sigmoid', padding='same')(x)

decoder = Model(input_rep,decoded)

autoencoder = Model(input_img, decoder(encoder(input_img)))
autoencoder.compile(optimizer='adam', loss='binary_crossentropy')

autoencoder.fit(x_train_lr, x_train,
                epochs=2,
                batch_size=128,
                shuffle=True,
                validation_data=(x_test_lr, x_test))

y_test_lr = autoencoder.predict(x_test_lr[0:5])
plotn(5,x_test_lr)
plotn(5,y_test_lr)

intermediate_dim = 512
latent_dim = 2
batch_size = 128

class Sampling(tf.keras.layers.Layer):
    """Samples a latent vector and registers its KL-divergence loss."""

    def call(self, inputs):
        z_mean, z_log_var = inputs
        epsilon = tf.random.normal(shape=tf.shape(z_mean))
        kl = -0.5 * tf.reduce_sum(
            1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var), axis=1
        )
        self.add_loss(tf.reduce_mean(kl))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

inputs = Input(shape=(784,), name="encoder_input")
h = Dense(intermediate_dim, activation="relu")(inputs)
z_mean = Dense(latent_dim, name="z_mean")(h)
z_log_var = Dense(latent_dim, name="z_log_var")(h)

z = Sampling(name="sampling")([z_mean, z_log_var])

encoder = Model(inputs, [z_mean, z_log_var, z], name="encoder")

latent_inputs = Input(shape=(latent_dim,), name="decoder_input")
x = Dense(intermediate_dim, activation="relu")(latent_inputs)
outputs = Dense(784, activation="sigmoid")(x)
decoder = Model(latent_inputs, outputs, name="decoder")

vae_outputs = decoder(encoder(inputs)[2])
vae = Model(inputs, vae_outputs, name="vae")

vae.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
    loss=tf.keras.losses.BinaryCrossentropy(),
)

x_train_flat = x_train.reshape((len(x_train), -1)).astype("float32")
x_test_flat = x_test.reshape((len(x_test), -1)).astype("float32")

vae.fit(
    x_train_flat,
    x_train_flat,
    shuffle=True,
    epochs=2,
    batch_size=batch_size,
    validation_data=(x_test_flat, x_test_flat),
)

y_test = vae.predict(x_test_flat[0:5])
plotn(5,x_test_flat)
plotn(5,y_test)

x_test_encoded = encoder.predict(x_test_flat)[0]
plt.figure(figsize=(6, 6))
plt.scatter(x_test_encoded[:, 0], x_test_encoded[:, 1], c=y_testclass)
plt.colorbar()
plt.show()

def plotsample(n):
  dx = np.linspace(-1,1,n)
  dy = np.linspace(-1,1,n)
  fig,ax = plt.subplots(n,n)
  for i,xi in enumerate(dx):
    for j,xj in enumerate(dy):
      res = decoder.predict(np.array([xi,xj]).reshape(-1,2))[0]
      ax[i,j].imshow(res.reshape(28,28))
      ax[i,j].axis('off')
  plt.show()

plotsample(10)