Updated 19/Nov/2021 by Yoshihisa Nitta  

AutoEncoder Analysis for MNIST dataset with Tensorflow 2 on Google Colab

MNIST データセットに対する AutoEncoder を Google Colab 上の Tensorflow 2を用いて解析する

In [1]:
#! pip install tensorflow==2.7.0
In [2]:
%tensorflow_version 2.x

import tensorflow as tf
print(tf.__version__)
2.7.0

Check the execution environment on Google Colab

Google Colab 上の実行環境を確認する

In [3]:
! nvidia-smi
! cat /proc/cpuinfo
! cat /etc/issue
! free -h
Mon Nov 22 08:51:34 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 495.44       Driver Version: 460.32.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla P100-PCIE...  Off  | 00000000:00:04.0 Off |                    0 |
| N/A   39C    P0    26W / 250W |      0MiB / 16280MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+
processor	: 0
vendor_id	: GenuineIntel
cpu family	: 6
model		: 79
model name	: Intel(R) Xeon(R) CPU @ 2.20GHz
stepping	: 0
microcode	: 0x1
cpu MHz		: 2199.998
cache size	: 56320 KB
physical id	: 0
siblings	: 2
core id		: 0
cpu cores	: 1
apicid		: 0
initial apicid	: 0
fpu		: yes
fpu_exception	: yes
cpuid level	: 13
wp		: yes
flags		: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx smap xsaveopt arat md_clear arch_capabilities
bugs		: cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs taa
bogomips	: 4399.99
clflush size	: 64
cache_alignment	: 64
address sizes	: 46 bits physical, 48 bits virtual
power management:

processor	: 1
vendor_id	: GenuineIntel
cpu family	: 6
model		: 79
model name	: Intel(R) Xeon(R) CPU @ 2.20GHz
stepping	: 0
microcode	: 0x1
cpu MHz		: 2199.998
cache size	: 56320 KB
physical id	: 0
siblings	: 2
core id		: 0
cpu cores	: 1
apicid		: 1
initial apicid	: 1
fpu		: yes
fpu_exception	: yes
cpuid level	: 13
wp		: yes
flags		: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx smap xsaveopt arat md_clear arch_capabilities
bugs		: cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs taa
bogomips	: 4399.99
clflush size	: 64
cache_alignment	: 64
address sizes	: 46 bits physical, 48 bits virtual
power management:

Ubuntu 18.04.5 LTS \n \l

              total        used        free      shared  buff/cache   available
Mem:            12G        741M          9G        1.2M        2.0G         11G
Swap:            0B          0B          0B

Mount Google Drive from Google Colab

Google Colab から GoogleDrive をマウントする

In [4]:
from google.colab import drive
drive.mount('/content/drive')
Mounted at /content/drive
In [5]:
! ls /content/drive
MyDrive  Shareddrives

Download the soure file from Google Drive or nw.tsuda.ac.jp

Basically, gdown from Google Drive. Download from nw.tsuda.ac.jp above only if the specifications of Google Drive change and you cannot download from Google Drive.

Google Drive または nw.tsuda.ac.jp からファイルをダウンロードする

基本的に Google Drive から gdown してください。 Google Drive の仕様が変わってダウンロードができない場合にのみ、nw.tsuda.ac.jp からダウンロードしてください。

In [6]:
# Download source file
nw_path = './nw'
! rm -rf {nw_path}
! mkdir -p {nw_path}

if True: # from Google Drive
    url_model = 'https://drive.google.com/uc?id=1ZDgWE7wmVwG_ZuQVUjuh_XHeIO-7Yn63'
    ! (cd {nw_path}; gdown {url_model})
else:    # from nw.tsuda.ac.jp
    URL_NW = 'https://nw.tsuda.ac.jp/lec/GoogleColab/pub'
    url_model = f'{URL_NW}/models/AutoEncoder.py'
    ! wget -nd {url_model} -P {nw_path}     # download to './nw/AutoEncoder.py'
Downloading...
From: https://drive.google.com/uc?id=1ZDgWE7wmVwG_ZuQVUjuh_XHeIO-7Yn63
To: /content/nw/AutoEncoder.py
100% 13.9k/13.9k [00:00<00:00, 25.7MB/s]
In [7]:
! cat {nw_path}/AutoEncoder.py
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

import os
import pickle
import datetime

class AutoEncoder():
    def __init__(self, 
                 input_dim,
                 encoder_conv_filters,
                 encoder_conv_kernel_size,
                 encoder_conv_strides,
                 decoder_conv_t_filters,
                 decoder_conv_t_kernel_size,
                 decoder_conv_t_strides,
                 z_dim,
                 use_batch_norm = False,
                 use_dropout = False,
                 epoch = 0
    ):
        self.name = 'autoencoder'
        self.input_dim = input_dim
        self.encoder_conv_filters = encoder_conv_filters
        self.encoder_conv_kernel_size = encoder_conv_kernel_size
        self.encoder_conv_strides = encoder_conv_strides
        self.decoder_conv_t_filters = decoder_conv_t_filters
        self.decoder_conv_t_kernel_size = decoder_conv_t_kernel_size
        self.decoder_conv_t_strides = decoder_conv_t_strides
        self.z_dim = z_dim
        
        self.use_batch_norm = use_batch_norm
        self.use_dropout = use_dropout

        self.epoch = epoch
        
        self.n_layers_encoder = len(encoder_conv_filters)
        self.n_layers_decoder = len(decoder_conv_t_filters)
        
        self._build()
 

    def _build(self):
        ### THE ENCODER
        encoder_input = tf.keras.layers.Input(shape=self.input_dim, name='encoder_input')
        x = encoder_input
    
        for i in range(self.n_layers_encoder):
            x = tf.keras.layers.Conv2D(
                filters = self.encoder_conv_filters[i],
                kernel_size = self.encoder_conv_kernel_size[i],
                strides = self.encoder_conv_strides[i],
                padding  = 'same',
                name = 'encoder_conv_' + str(i)
            )(x)
            x = tf.keras.layers.LeakyReLU()(x)
            if self.use_batch_norm:
                x = tf.keras.layers.BatchNormalization()(x)
            if self.use_dropout:
                x = tf.keras.layers.Dropout(rate = 0.25)(x)
              
        shape_before_flattening = tf.keras.backend.int_shape(x)[1:] # shape for 1 data
        
        x = tf.keras.layers.Flatten()(x)
        encoder_output = tf.keras.layers.Dense(self.z_dim, name='encoder_output')(x)
        
        self.encoder = tf.keras.models.Model(encoder_input, encoder_output)
        
        ### THE DECODER
        decoder_input = tf.keras.layers.Input(shape=(self.z_dim,), name='decoder_input')
        x = tf.keras.layers.Dense(np.prod(shape_before_flattening))(decoder_input)
        x = tf.keras.layers.Reshape(shape_before_flattening)(x)
        
        for i in range(self.n_layers_decoder):
            x =   tf.keras.layers.Conv2DTranspose(
                filters = self.decoder_conv_t_filters[i],
                kernel_size = self.decoder_conv_t_kernel_size[i],
                strides = self.decoder_conv_t_strides[i],
                padding = 'same',
                name = 'decoder_conv_t_' + str(i)
            )(x)
          
            if i < self.n_layers_decoder - 1:
                x = tf.keras.layers.LeakyReLU()(x)
                if self.use_batch_norm:
                    x = tf.keras.layers.BatchNormalization()(x)
                if self.use_dropout:
                    x = tf.keras.layers.Dropout(rate=0.25)(x)
            else:
                x = tf.keras.layers.Activation('sigmoid')(x)
              
        decoder_output = x
        self.decoder = tf.keras.models.Model(decoder_input, decoder_output)
        
        ### THE FULL AUTOENCODER
        model_input = encoder_input
        model_output = self.decoder(encoder_output)
        
        self.model = tf.keras.models.Model(model_input, model_output)


    def save(self, folder):
        self.save_params(os.path.join(folder, 'params.pkl'))
        self.save_weights(os.path.join(folder, 'weights/weights.h5'))


    @staticmethod
    def load(folder, epoch=None):   # AutoEncoder.load(folder)
        params = AutoEncoder.load_params(os.path.join(folder, 'params.pkl'))
        AE = AutoEncoder(*params)
        if epoch is None:
            AE.model.load_weights(os.path.join(folder, 'weights/weights.h5'))
        else:
            AE.model.load_weights(os.path.join(folder, f'weights/weights_{epoch-1}.h5'))
            AE.epoch = epoch

        return AE


    def save_params(self, filepath):
        dpath, fname = os.path.split(filepath)
        if dpath != '' and not os.path.exists(dpath):
            os.makedirs(dpath)
        with open(filepath, 'wb') as f:
            pickle.dump([
                self.input_dim,
                self.encoder_conv_filters,
                self.encoder_conv_kernel_size,
                self.encoder_conv_strides,
                self.decoder_conv_t_filters,
                self.decoder_conv_t_kernel_size,
                self.decoder_conv_t_strides,
                self.z_dim,
                self.use_batch_norm,
                self.use_dropout,
                self.epoch
            ], f)


    @staticmethod
    def load_params(filepath):
        with open(filepath, 'rb') as f:
            params = pickle.load(f)
        return params


    def save_weights(self, filepath):
        dpath, fname = os.path.split(filepath)
        if dpath != '' and not os.path.exists(dpath):
            os.makedirs(dpath)
        self.model.save_weights(filepath)
        
        
    def load_weights(self, filepath):
        self.model.load_weights(filepath)


    def save_images(self, imgs, filepath):
        z_points = self.encoder.predict(imgs)
        reconst_imgs = self.decoder.predict(z_points)
        txts = [ f'{p[0]:.3f}, {p[1]:.3f}' for p in z_points ]
        AutoEncoder.showImages(imgs, reconst_imgs, txts, 1.4, 1.4, 0.5, filepath)
      

    @staticmethod
    def r_loss(y_true, y_pred):
        return tf.keras.backend.mean(tf.keras.backend.square(y_true - y_pred), axis=[1,2,3])


    def compile(self, learning_rate):
        self.learning_rate = learning_rate
        optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
        self.model.compile(optimizer=optimizer, loss = AutoEncoder.r_loss)

        
    def train_with_fit(self,
               x_train,
               y_train,
               batch_size,
               epochs,
               run_folder='run/',
               validation_data=None
    ):
        history= self.model.fit(
            x_train,
            y_train,
            batch_size = batch_size,
            shuffle = True,
            initial_epoch = self.epoch,
            epochs = epochs,
            validation_data = validation_data
        )
        if self.epoch < epochs:
            self.epoch = epochs

        if run_folder != None:
            self.save(run_folder)
            self.save_weights(os.path.join(run_folder,f'weights/weights_{self.epoch-1}.h5'))
            #idxs = np.random.choice(len(x_train), 10)
            #self.save_images(x_train[idxs], os.path.join(run_folder, f'images/image_{self.epoch-1}.png'))

        return history
        
        
    def train(self,
               x_train,
               y_train,
               batch_size = 32,
               epochs = 10,
               shuffle=False,
               run_folder='run/',
               optimizer=None,
               save_epoch_interval=100,
               validation_data = None
    ):
        start_time = datetime.datetime.now()
        steps = x_train.shape[0] // batch_size

        losses = []
        val_losses = []

        for epoch in range(self.epoch, epochs):
            epoch_loss = 0
            indices = tf.range(x_train.shape[0], dtype=tf.int32)
            if shuffle:
                indices = tf.random.shuffle(indices)
            x_ = x_train[indices]
            y_ = y_train[indices]
            
            for step in range(steps):
                start = batch_size * step
                end = start + batch_size

                with tf.GradientTape() as tape:
                    outputs = self.model(x_[start:end])
                    tmp_loss = AutoEncoder.r_loss(y_[start:end], outputs)

                grads = tape.gradient(tmp_loss, self.model.trainable_variables)
                optimizer.apply_gradients(zip(grads, self.model.trainable_variables))

            epoch_loss = np.mean(tmp_loss)
            losses.append(epoch_loss)

            val_str = ''
            if validation_data != None:
                x_val, y_val = validation_data
                outputs_val = self.model(x_val)
                val_loss = np.mean(AutoEncoder.r_loss(y_val, outputs_val))
                val_str = f'val loss: {val_loss:.4f}  '
                val_losses.append(val_loss)


            if (epoch+1) % save_epoch_interval == 0 and run_folder != None:
                self.save(run_folder)
                self.save_weights(os.path.join(run_folder,f'weights/weights_{self.epoch}.h5'))
                #idxs = np.random.choice(len(x_train), 10)
                #self.save_images(x_train[idxs], os.path.join(run_folder, f'images/image_{self.epoch}.png'))

            elapsed_time = datetime.datetime.now() - start_time
            print(f'{epoch+1}/{epochs} {steps} loss: {epoch_loss:.4f}  {val_str}{elapsed_time}')

            self.epoch += 1

        if run_folder != None:
            self.save(run_folder)
            self.save_weights(os.path.join(run_folder,f'weights/weights_{self.epoch-1}.h5'))
            #idxs = np.random.choice(len(x_train), 10)
            #self.save_images(x_train[idxs], os.path.join(run_folder, f'images/image_{self.epoch-1}.png'))

        return losses, val_losses

    @staticmethod
    @tf.function
    def compute_loss_and_grads(model,x,y):
        with tf.GradientTape() as tape:
            outputs = model(x)
            tmp_loss = AutoEncoder.r_loss(y,outputs)
        grads = tape.gradient(tmp_loss, model.trainable_variables)
        return tmp_loss, grads


    def train_tf(self,
               x_train,
               y_train,
               batch_size = 32,
               epochs = 10,
               shuffle=False,
               run_folder='run/',
               optimizer=None,
               save_epoch_interval=100,
               validation_data = None
    ):
        start_time = datetime.datetime.now()
        steps = x_train.shape[0] // batch_size

        losses = []
        val_losses = []

        for epoch in range(self.epoch, epochs):
            epoch_loss = 0
            indices = tf.range(x_train.shape[0], dtype=tf.int32)
            if shuffle:
                indices = tf.random.shuffle(indices)
            x_ = x_train[indices]
            y_ = y_train[indices]

            step_losses = []
            for step in range(steps):
                start = batch_size * step
                end = start + batch_size

                tmp_loss, grads = AutoEncoder.compute_loss_and_grads(self.model, x_[start:end], y_[start:end])
                optimizer.apply_gradients(zip(grads, self.model.trainable_variables))

                step_losses.append(np.mean(tmp_loss))

            epoch_loss = np.mean(step_losses)
            losses.append(epoch_loss)

            val_str = ''
            if validation_data != None:
                x_val, y_val = validation_data
                outputs_val = self.model(x_val)
                val_loss = np.mean(AutoEncoder.r_loss(y_val, outputs_val))
                val_str = f'val loss: {val_loss:.4f}  '
                val_losses.append(val_loss)


            if (epoch+1) % save_epoch_interval == 0 and run_folder != None:
                self.save(run_folder)
                self.save_weights(os.path.join(run_folder,f'weights/weights_{self.epoch}.h5'))
                #idxs = np.random.choice(len(x_train), 10)
                #self.save_images(x_train[idxs], os.path.join(run_folder, f'images/image_{self.epoch}.png'))

            elapsed_time = datetime.datetime.now() - start_time
            print(f'{epoch+1}/{epochs} {steps} loss: {epoch_loss:.4f}  {val_str}{elapsed_time}')

            self.epoch += 1

        if run_folder != None:
            self.save(run_folder)
            self.save_weights(os.path.join(run_folder,f'weights/weights_{self.epoch-1}.h5'))
            #idxs = np.random.choice(len(x_train), 10)
            #self.save_images(x_train[idxs], os.path.join(run_folder, f'images/image_{self.epoch-1}.png'))

        return losses, val_losses


    @staticmethod
    def showImages(imgs1, imgs2, txts, w, h, vskip=0.5, filepath=None):
        n = len(imgs1)
        fig, ax = plt.subplots(2, n, figsize=(w * n, (2+vskip) * h))
        for i in range(n):
            if n == 1:
                axis = ax[0]
            else:
                axis = ax[0][i]
            img = imgs1[i].squeeze()
            axis.imshow(img, cmap='gray_r')
            axis.axis('off')

            axis.text(0.5, -0.35, txts[i], fontsize=10, ha='center', transform=axis.transAxes)

            if n == 1:
                axis = ax[1]
            else:
                axis = ax[1][i]
            img2 = imgs2[i].squeeze()
            axis.imshow(img2, cmap='gray_r')
            axis.axis('off')

        if not filepath is None:
            dpath, fname = os.path.split(filepath)
            if dpath != '' and not os.path.exists(dpath):
                os.makedirs(dpath)
            fig.savefig(filepath, dpi=600)
            plt.close()
        else:
            plt.show()

    @staticmethod
    def plot_history(vals, labels):
        colors = ['red', 'blue', 'green', 'orange', 'black']
        n = len(vals)
        fig, ax = plt.subplots(1, 1, figsize=(9,4))
        for i in range(n):
            ax.plot(vals[i], c=colors[i], label=labels[i])
        ax.legend(loc='upper right')
        ax.set_xlabel('epochs')
        # ax[0].set_ylabel('loss')
        
        plt.show()

Preparing the MNIST datasets

MNIST データセットを用意する

In [8]:
import tensorflow as tf
import numpy as np
In [9]:
# MNIST datasets
(x_train_raw, y_train_raw), (x_test_raw, y_test_raw) = tf.keras.datasets.mnist.load_data()
print(x_train_raw.shape)
print(y_train_raw.shape)
print(x_test_raw.shape)
print(y_test_raw.shape)
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step
11501568/11490434 [==============================] - 0s 0us/step
(60000, 28, 28)
(60000,)
(10000, 28, 28)
(10000,)
In [10]:
x_train = x_train_raw.reshape(x_train_raw.shape+(1,)).astype('float32') / 255.0
x_test = x_test_raw.reshape(x_test_raw.shape+(1,)).astype('float32') / 255.0
print(x_train.shape)
print(x_test.shape)
(60000, 28, 28, 1)
(10000, 28, 28, 1)

Load the Neural Network Model trained before.

Load the Neural Network Model trained in the AE_MNIST_Train.ipynb. The definition of AutoEncoder class is downloaded from nw.tsuda.ac.jp.

学習済みのニューラルネットワーク・モデルをロードする

AE_MNIST_Train.ipynb で学習したニューラルネットワーク・モデルをロードする。 AutoEncoder クラスの定義は nw.tsuda.ac.jp からダウンロードした。

In [11]:
save_path1 = '/content/drive/MyDrive/ColabRun/AE01'
save_path2 = '/content/drive/MyDrive/ColabRun/AE02'
save_path3 = '/content/drive/MyDrive/ColabRun/AE03'
In [12]:
from nw.AutoEncoder import AutoEncoder
AE = AutoEncoder.load(save_path3)
In [13]:
print(AE.epoch)
200
In [14]:
# Load the AutoEncoder object of the specified epoch.
# 保存したAutoEncoderクラスのうち、指定したepochのオブジェクトをロードする。

AE_young = AutoEncoder.load(save_path3, 3)

print(AE_young.epoch)
3

Display the distribution of points in the latent space.

Encode 10000 images of x_test to generated 2D coordinates and draw as points. The label of the image is expressed by the color of the dots.

潜在空間における点の分布を表示する

x_test の10000枚の画像をエンコードして2次元座標を生成し、点として描画する。各画像に描かれている数字は、点の色で表現する。

In [15]:
n_to_show = len(x_test)
example_idx = np.random.choice(range(len(x_test)), n_to_show)
example_images = x_test[example_idx]
example_labels = y_test_raw[example_idx]

z_points = AE.encoder.predict(example_images)

min_x = min(z_points[:, 0])
max_x = max(z_points[:, 0])
min_y = min(z_points[:, 1])
max_y = max(z_points[:, 1])
In [16]:
%matplotlib inline
import matplotlib.pyplot as plt

fig, ax = plt.subplots(1, 1, figsize=(12, 12))
map = ax.scatter(z_points[:, 0], z_points[:, 1], c=example_labels, cmap='rainbow', alpha=0.5, s=2)

plt.colorbar(map)

plt.show()
In [17]:
# Display 30 images as points in the latent space.
# 30 枚の画像を潜在空間の点として表示する。
import numpy as np

table_row = 10 
table_line = 3

x = np.random.uniform(min_x, max_x, size=table_line * table_row)
y = np.random.uniform(min_y, max_y, size=table_line * table_row)
z_grid = np.array(list(zip(x,y)))    # (x, y) : 2D coordinates
reconst = AE.decoder.predict(z_grid)
In [18]:
%matplotlib inline
import matplotlib.pyplot as plt

fig, ax = plt.subplots(1, 1, figsize=(12, 12))
map = ax.scatter(z_points[:, 0], z_points[:, 1], c=example_labels, cmap='rainbow', alpha=0.5, s=2)

ax.scatter(z_grid[:,0], z_grid[:,1], c='red', alpha=1, s=20)

for i in range(len(z_grid)):
    ax.text(z_grid[i][0], z_grid[i][1], str(i))

plt.colorbar(map)

plt.show()
In [19]:
# Display the 30 original images selected as samples.
# サンプルとして選んだ30枚の元画像を表示する
%matplotlib inline
import matplotlib.pyplot as plt

VSKIP=0.5   # vertical space between subplots

fig, ax = plt.subplots(table_line, table_row, figsize=(2.8 * table_row, 2.8 * table_line * (1+VSKIP)))
plt.subplots_adjust(hspace = VSKIP)
                       
for y in range(table_line):
    for x in range(table_row):
        idx = table_row * y + x
        img = reconst[idx].squeeze()
        ax[y][x].imshow(img, cmap='gray')
        ax[y][x].text(0.5, -0.35, f'{idx} ({z_grid[idx][0]:.2f}, {z_grid[idx][1]:.2f})', fontsize=16, ha='center', transform=ax[y][x].transAxes)
        ax[y][x].axis('off')
        
plt.show()

Divide the latent space into grids and find out what kind of image is generated (decoded) from each coordinates.

Generate images from points on 20x20 grid. The generated images are displayed as a table with 20 rows and 20 columns.

潜在空間をグリッドに区切って、各座標からどのような画像が生成(デコード)されるか調べる

20×20 のグリッドから、画像を生成する。 生成した画像は 20 行 20列の表として表示する。

In [20]:
import numpy as np

n_grid = 20

x = np.linspace(min_x, max_x, n_grid)
y = np.linspace(min_y, max_y, n_grid)

xv, yv = np.meshgrid(x, y)
xv = xv.flatten()
yv = yv.flatten()
z_grid2 = np.array(list(zip(xv, yv)))

reconst2 = AE.decoder.predict(z_grid2)
In [21]:
%matplotlib inline
import matplotlib.pyplot as plt

fig, ax = plt.subplots(n_grid, n_grid, figsize=(n_grid, n_grid))
for i in range(len(reconst2)):
    img = reconst2[i].squeeze()
    line = n_grid - 1 - i // n_grid
    row = i % n_grid
    ax[line][row].imshow(img, cmap='gray')
    ax[line][row].axis('off')
    
plt.show()