Add save/load functionality to VAE model

Former-commit-id: bdfef613

VAE/01-VAE-with-MNIST.ipynb

+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Variational AutoEncoder (VAE) with MNIST\n",
+    "========================================\n",
+    "---\n",
+    "Formation Introduction au Deep Learning  (FIDLE) - S. Arias, E. Maldonado, JL. Parouty - CNRS/SARI/DEVLOG - 2020  \n",
+    "\n",
+    "## Episode 1 - Train a model\n",
+    " - Defining a VAE model\n",
+    " - Build the model\n",
+    " - Train it\n",
+    " - Follow the learning process with Tensorboard\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 1 - Init python stuff"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 65,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "FIDLE 2020 - Variational AutoEncoder (VAE)\n",
+      "TensorFlow version : 2.0.0\n"
+     ]
+    }
+   ],
+   "source": [
+    "import numpy as np\n",
+    "import tensorflow as tf\n",
+    "import tensorflow.keras.datasets.mnist as mnist\n",
+    "import sys, importlib\n",
+    "\n",
+    "import modules.vae\n",
+    "importlib.reload(modules.vae)\n",
+    "\n",
+    "print('FIDLE 2020 - Variational AutoEncoder (VAE)')\n",
+    "print('TensorFlow version :',tf.__version__)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 2 - Get data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 66,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Dataset loaded.\n",
+      "x_train shape : (60000, 28, 28, 1)\n",
+      "x_test_shape  : (10000, 28, 28, 1)\n"
+     ]
+    }
+   ],
+   "source": [
+    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
+    "\n",
+    "x_train = x_train.astype('float32') / 255.\n",
+    "x_train = np.expand_dims(x_train, axis=3)\n",
+    "x_test  = x_test.astype('float32') / 255.\n",
+    "x_test  = np.expand_dims(x_test, axis=3)\n",
+    "print('Dataset loaded.')\n",
+    "print(f'x_train shape : {x_train.shape}\\nx_test_shape  : {x_test.shape}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 3 - Get VAE model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 67,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Model initialized.\n",
+      "Outputs will be in :  ./run/004\n"
+     ]
+    }
+   ],
+   "source": [
+    "tag = '004'\n",
+    "\n",
+    "input_shape = (28,28,1)\n",
+    "z_dim       = 2\n",
+    "verbose     = 0\n",
+    "\n",
+    "encoder= [ {'type':'Conv2D', 'filters':32, 'kernel_size':(3,3), 'strides':1, 'padding':'same', 'activation':'relu'},\n",
+    "           {'type':'Conv2D', 'filters':64, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
+    "           {'type':'Conv2D', 'filters':64, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
+    "           {'type':'Conv2D', 'filters':64, 'kernel_size':(3,3), 'strides':1, 'padding':'same', 'activation':'relu'}\n",
+    "         ]\n",
+    "\n",
+    "decoder= [ {'type':'Conv2DT', 'filters':64, 'kernel_size':(3,3), 'strides':1, 'padding':'same', 'activation':'relu'},\n",
+    "           {'type':'Conv2DT', 'filters':64, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
+    "           {'type':'Conv2DT', 'filters':32, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
+    "           {'type':'Conv2DT', 'filters':1,  'kernel_size':(3,3), 'strides':1, 'padding':'same', 'activation':'sigmoid'}\n",
+    "         ]\n",
+    "\n",
+    "vae = modules.vae.VariationalAutoencoder(input_shape    = input_shape, \n",
+    "                                         encoder_layers = encoder, \n",
+    "                                         decoder_layers = decoder,\n",
+    "                                         z_dim          = z_dim, \n",
+    "                                         verbose        = verbose,\n",
+    "                                         run_tag        = tag)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 68,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Config saved in : ./run/004/models/vae_config.json\n",
+      "Model saved in  : ./run/004/models/model.h5\n"
+     ]
+    }
+   ],
+   "source": [
+    "vae.save()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 69,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dict_keys(['input_shape', 'encoder_layers', 'decoder_layers', 'z_dim', 'run_tag', 'verbose'])\n",
+      "Model initialized.\n",
+      "Outputs will be in :  ./run/004\n",
+      "Weights loaded from : ./run/004/models/model.h5\n"
+     ]
+    }
+   ],
+   "source": [
+    "vae=modules.vae.VariationalAutoencoder.load('004')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 4 - Compile it"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 70,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Compiled.\n",
+      "Optimizer is Adam with learning_rate=0.0005\n"
+     ]
+    }
+   ],
+   "source": [
+    "learning_rate = 0.0005\n",
+    "r_loss_factor = 1000\n",
+    "\n",
+    "vae.compile(learning_rate, r_loss_factor)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 5 - Train"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_size        = 100\n",
+    "epochs            = 100\n",
+    "image_periodicity = 1      # for each epoch\n",
+    "chkpt_periodicity = 2      # for each epoch\n",
+    "initial_epoch     = 0\n",
+    "dataset_size      = 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vae.train(x_train,\n",
+    "          x_test,\n",
+    "          batch_size        = batch_size, \n",
+    "          epochs            = epochs,\n",
+    "          image_periodicity = image_periodicity,\n",
+    "          chkpt_periodicity = chkpt_periodicity,\n",
+    "          initial_epoch     = initial_epoch,\n",
+    "          dataset_size      = dataset_size\n",
+    "         )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vae."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

VAE/02-VAE-with-MNIST-post.ipynb

+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Variational AutoEncoder (VAE) with MNIST\n",
+    "========================================\n",
+    "---\n",
+    "Formation Introduction au Deep Learning  (FIDLE) - S. Arias, E. Maldonado, JL. Parouty - CNRS/SARI/DEVLOG - 2020  \n",
+    "\n",
+    "## Episode 2 - Analyse our trained model\n",
+    " - Defining a VAE model\n",
+    " - Build the model\n",
+    " - Train it\n",
+    " - Follow the learning process with Tensorboard\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 1 - Init python stuff"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "FIDLE 2020 - Variational AutoEncoder (VAE)\n",
+      "TensorFlow version : 2.0.0\n"
+     ]
+    }
+   ],
+   "source": [
+    "import numpy as np\n",
+    "import tensorflow as tf\n",
+    "import tensorflow.keras.datasets.mnist as mnist\n",
+    "import sys, importlib\n",
+    "\n",
+    "import modules.vae\n",
+    "importlib.reload(modules.vae)\n",
+    "\n",
+    "print('FIDLE 2020 - Variational AutoEncoder (VAE)')\n",
+    "print('TensorFlow version :',tf.__version__)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 2 - Get data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Dataset loaded.\n",
+      "x_train shape : (60000, 28, 28, 1)\n",
+      "x_test_shape  : (10000, 28, 28, 1)\n"
+     ]
+    }
+   ],
+   "source": [
+    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
+    "\n",
+    "x_train = x_train.astype('float32') / 255.\n",
+    "x_train = np.expand_dims(x_train, axis=3)\n",
+    "x_test  = x_test.astype('float32') / 255.\n",
+    "x_test  = np.expand_dims(x_test, axis=3)\n",
+    "print('Dataset loaded.')\n",
+    "print(f'x_train shape : {x_train.shape}\\nx_test_shape  : {x_test.shape}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 3 - Load best model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vae\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "end_time  = time.time()\n",
+    "dt  = end_time-start_time\n",
+    "dth = str(datetime.timedelta(seconds=dt))\n",
+    "print(f'\\nTrain duration : {dt:.2f} sec. - {dth:}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 4 - Compile it"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "learning_rate = 0.0005\n",
+    "r_loss_factor = 1000\n",
+    "\n",
+    "vae.compile(learning_rate, r_loss_factor)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 5 - Train"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_size        = 100\n",
+    "epochs            = 200\n",
+    "image_periodicity = 1      # for each epoch\n",
+    "chkpt_periodicity = 2      # for each epoch\n",
+    "initial_epoch     = 0\n",
+    "dataset_size      = 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vae.train(x_train,\n",
+    "          x_test,\n",
+    "          batch_size        = batch_size, \n",
+    "          epochs            = epochs,\n",
+    "          image_periodicity = image_periodicity,\n",
+    "          chkpt_periodicity = chkpt_periodicity,\n",
+    "          initial_epoch     = initial_epoch,\n",
+    "          dataset_size      = dataset_size,\n",
+    "          lr_decay          = 1\n",
+    "         )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

VAE/modules/vae.py

@@ -13,20 +13,21 @@ from tensorflow.keras.utils import plot_model
 import tensorflow.keras.datasets.imdb as imdb
 
 import modules.callbacks
-import os
+import os, json, time, datetime
 
 
 
 class VariationalAutoencoder():
 
     
-    def __init__(self, input_shape=None, encoder_layers=None, decoder_layers=None, z_dim=None, run_tag='default', verbose=0):
-        
+    def __init__(self, input_shape=None, encoder_layers=None, decoder_layers=None, z_dim=None, run_tag='000', verbose=0):
+               
         self.name           = 'Variational AutoEncoder'
-        self.input_shape    = input_shape
+        self.input_shape    = list(input_shape)
         self.encoder_layers = encoder_layers
         self.decoder_layers = decoder_layers
         self.z_dim          = z_dim
+        self.run_tag        = str(run_tag)
         self.verbose        = verbose
         self.run_directory  = f'./run/{run_tag}'
         
@@ -42,13 +43,14 @@ class VariationalAutoencoder():
         
         # ---- Add next layers
         i=1
-        for params in encoder_layers:
-            t=params['type']
-            params.pop('type')
-            if t=='Conv2D':
-                layer = Conv2D(**params, name=f"Layer_{i}")
-            if t=='Dropout':
-                layer = Dropout(**params)
+        for l_config in encoder_layers:
+            l_type   = l_config['type']
+            l_params = l_config.copy()
+            l_params.pop('type')
+            if l_type=='Conv2D':
+                layer = Conv2D(**l_params)
+            if l_type=='Dropout':
+                layer = Dropout(**l_params)
             x = layer(x)
             i+=1
             
@@ -83,13 +85,14 @@ class VariationalAutoencoder():
 
         # ---- Add next layers
         i=1
-        for params in decoder_layers:
-            t=params['type']
-            params.pop('type')
-            if t=='Conv2DT':
-                layer = Conv2DTranspose(**params, name=f"Layer_{i}")
-            if t=='Dropout':
-                layer = Dropout(**params)
+        for l_config in decoder_layers:
+            l_type   = l_config['type']
+            l_params = l_config.copy()
+            l_params.pop('type')
+            if l_type=='Conv2DT':
+                layer = Conv2DTranspose(**l_params)
+            if l_type=='Dropout':
+                layer = Dropout(**l_params)
             x = layer(x)
             i+=1
 
@@ -140,6 +143,8 @@ class VariationalAutoencoder():
                            loss = vae_loss,
                            metrics = [vae_r_loss, vae_kl_loss], 
                            experimental_run_tf_function=False)
+        print('Compiled.')
+        print(f'Optimizer is Adam with learning_rate={learning_rate:}')
     
     
     def train(self, 
@@ -165,7 +170,7 @@ class VariationalAutoencoder():
         callbacks_images = modules.callbacks.ImagesCallback(initial_epoch, image_periodicity, self)
         
         # ---- Callback : Learning rate scheduler
-        lr_sched = modules.callbacks.step_decay_schedule(initial_lr=self.learning_rate, decay_factor=lr_decay, step_size=1)
+        #lr_sched = modules.callbacks.step_decay_schedule(initial_lr=self.learning_rate, decay_factor=lr_decay, step_size=1)
         
         # ---- Callback : Checkpoint
         filename = self.run_directory+"/models/model-{epoch:03d}-{loss:.2f}.h5"
@@ -179,17 +184,23 @@ class VariationalAutoencoder():
         dirname = self.run_directory+"/logs"
         callback_tensorboard = TensorBoard(log_dir=dirname, histogram_freq=1)
 
-        callbacks_list = [callbacks_images, callback_chkpts, callback_bestmodel, callback_tensorboard, lr_sched]
-
-        self.model.fit(x_train[:n_train], x_train[:n_train],
-                       batch_size = batch_size,
-                       shuffle = True,
-                       epochs = epochs,
-                       initial_epoch = initial_epoch,
-                       callbacks = callbacks_list,
-                       validation_data = (x_test[:n_test], x_test[:n_test])
-                        )
-        
+        callbacks_list = [callbacks_images, callback_chkpts, callback_bestmodel, callback_tensorboard]
+
+        # ---- Let's go...
+        start_time   = time.time()
+        self.history = self.model.fit(x_train[:n_train], x_train[:n_train],
+                                      batch_size = batch_size,
+                                      shuffle = True,
+                                      epochs = epochs,
+                                      initial_epoch = initial_epoch,
+                                      callbacks = callbacks_list,
+                                      validation_data = (x_test[:n_test], x_test[:n_test])
+                                      )
+        end_time  = time.time()
+        dt  = end_time-start_time
+        dth = str(datetime.timedelta(seconds=int(dt)))
+        self.duration = dt
+        print(f'\nTrain duration : {dt:.2f} sec. - {dth:}')
         
     def plot_model(self):
         d=self.run_directory+'/figs'
@@ -198,3 +209,39 @@ class VariationalAutoencoder():
         plot_model(self.decoder, to_file=f'{d}/decoder.png', show_shapes = True, show_layer_names = True)
 
         
+    def save(self,config='vae_config.json', model='model.h5'):
+        # ---- Save config in json
+        if config!=None:
+            to_save  = ['input_shape', 'encoder_layers', 'decoder_layers', 'z_dim', 'run_tag', 'verbose']
+            data     = { i:self.__dict__[i] for i in to_save }
+            filename = self.run_directory+'/models/'+config
+            with open(filename, 'w') as outfile:
+                json.dump(data, outfile)
+            print(f'Config saved in : {filename}')
+        # ---- Save model
+        if model!=None:
+            filename = self.run_directory+'/models/'+model
+            self.model.save(filename)
+            print(f'Model saved in  : {filename}')
+
+            
+    def load_weights(self,model='model.h5'):
+        filename = self.run_directory+'/models/'+model
+        self.model.load_weights(filename)
+        print(f'Weights loaded from : {filename}')
+    
+            
+    @classmethod
+    def load(cls, run_tag='000', config='vae_config.json', model='model.h5'):
+        # ---- Instantiate a new vae
+        filename = f'./run/{run_tag}/models/{config}'
+        with open(filename, 'r') as infile:
+            params=json.load(infile)
+            print(params.keys())
+#             vae=cls( params['input_shape'], params['encoder_layers'], params['decoder_layers'], params['z_dim'], '004', 0)
+            vae=cls( **params)
+        # ---- model==None, just return it
+        if model==None: return vae
+        # ---- model!=None, get weight
+        vae.load_weights(model)
+        return vae
\ No newline at end of file