{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save & Restore a Model\n",
    "\n",
    "Save and Restore a model using TensorFlow v2. In this example, we will go over both low and high-level approaches: \n",
    "- Low-level: TF Checkpoint.\n",
    "- High-level: TF Module/Model saver.\n",
    "\n",
    "This example is using the MNIST database of handwritten digits as toy dataset\n",
    "(http://yann.lecun.com/exdb/mnist/).\n",
    "\n",
    "- Author: Aymeric Damien\n",
    "- Project: https://github.com/aymericdamien/TensorFlow-Examples/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MNIST dataset parameters.\n",
    "num_classes = 10 # 0 to 9 digits\n",
    "num_features = 784 # 28*28\n",
    "\n",
    "# Training parameters.\n",
    "learning_rate = 0.01\n",
    "training_steps = 1000\n",
    "batch_size = 256\n",
    "display_step = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare MNIST data.\n",
    "from tensorflow.keras.datasets import mnist\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "# Convert to float32.\n",
    "x_train, x_test = np.array(x_train, np.float32), np.array(x_test, np.float32)\n",
    "# Flatten images to 1-D vector of 784 features (28*28).\n",
    "x_train, x_test = x_train.reshape([-1, num_features]), x_test.reshape([-1, num_features])\n",
    "# Normalize images value from [0, 255] to [0, 1].\n",
    "x_train, x_test = x_train / 255., x_test / 255."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use tf.data API to shuffle and batch data.\n",
    "train_data = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
    "train_data = train_data.repeat().shuffle(5000).batch(batch_size).prefetch(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) TF Checkpoint\n",
    "\n",
    "Basic logistic regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Weight of shape [784, 10], the 28*28 image features, and total number of classes.\n",
    "W = tf.Variable(tf.random.normal([num_features, num_classes]), name=\"weight\")\n",
    "# Bias of shape [10], the total number of classes.\n",
    "b = tf.Variable(tf.zeros([num_classes]), name=\"bias\")\n",
    "\n",
    "# Logistic regression (Wx + b).\n",
    "def logistic_regression(x):\n",
    "    # Apply softmax to normalize the logits to a probability distribution.\n",
    "    return tf.nn.softmax(tf.matmul(x, W) + b)\n",
    "\n",
    "# Cross-Entropy loss function.\n",
    "def cross_entropy(y_pred, y_true):\n",
    "    # Encode label to a one hot vector.\n",
    "    y_true = tf.one_hot(y_true, depth=num_classes)\n",
    "    # Clip prediction values to avoid log(0) error.\n",
    "    y_pred = tf.clip_by_value(y_pred, 1e-9, 1.)\n",
    "    # Compute cross-entropy.\n",
    "    return tf.reduce_mean(-tf.reduce_sum(y_true * tf.math.log(y_pred)))\n",
    "\n",
    "# Accuracy metric.\n",
    "def accuracy(y_pred, y_true):\n",
    "    # Predicted class is the index of highest score in prediction vector (i.e. argmax).\n",
    "    correct_prediction = tf.equal(tf.argmax(y_pred, 1), tf.cast(y_true, tf.int64))\n",
    "    return tf.reduce_mean(tf.cast(correct_prediction, tf.float32))\n",
    "\n",
    "# Adam optimizer.\n",
    "optimizer = tf.optimizers.Adam(learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimization process. \n",
    "def run_optimization(x, y):\n",
    "    # Wrap computation inside a GradientTape for automatic differentiation.\n",
    "    with tf.GradientTape() as g:\n",
    "        pred = logistic_regression(x)\n",
    "        loss = cross_entropy(pred, y)\n",
    "\n",
    "        # Compute gradients.\n",
    "        gradients = g.gradient(loss, [W, b])\n",
    "\n",
    "        # Update W and b following gradients.\n",
    "        optimizer.apply_gradients(zip(gradients, [W, b]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step: 50, loss: 535.380981, accuracy: 0.656250\n",
      "step: 100, loss: 354.681152, accuracy: 0.765625\n",
      "step: 150, loss: 225.300934, accuracy: 0.785156\n",
      "step: 200, loss: 163.948761, accuracy: 0.859375\n",
      "step: 250, loss: 129.653534, accuracy: 0.878906\n",
      "step: 300, loss: 170.743576, accuracy: 0.859375\n",
      "step: 350, loss: 97.912575, accuracy: 0.910156\n",
      "step: 400, loss: 144.119141, accuracy: 0.906250\n",
      "step: 450, loss: 164.991943, accuracy: 0.875000\n",
      "step: 500, loss: 145.191666, accuracy: 0.871094\n",
      "step: 550, loss: 82.272644, accuracy: 0.925781\n",
      "step: 600, loss: 149.180237, accuracy: 0.878906\n",
      "step: 650, loss: 127.171280, accuracy: 0.871094\n",
      "step: 700, loss: 116.045761, accuracy: 0.910156\n",
      "step: 750, loss: 92.582680, accuracy: 0.906250\n",
      "step: 800, loss: 108.238007, accuracy: 0.894531\n",
      "step: 850, loss: 92.755638, accuracy: 0.894531\n",
      "step: 900, loss: 69.131119, accuracy: 0.902344\n",
      "step: 950, loss: 67.176285, accuracy: 0.921875\n",
      "step: 1000, loss: 104.205658, accuracy: 0.890625\n"
     ]
    }
   ],
   "source": [
    "# Run training for the given number of steps.\n",
    "for step, (batch_x, batch_y) in enumerate(train_data.take(training_steps), 1):\n",
    "    # Run the optimization to update W and b values.\n",
    "    run_optimization(batch_x, batch_y)\n",
    "    \n",
    "    if step % display_step == 0:\n",
    "        pred = logistic_regression(batch_x)\n",
    "        loss = cross_entropy(pred, batch_y)\n",
    "        acc = accuracy(pred, batch_y)\n",
    "        print(\"step: %i, loss: %f, accuracy: %f\" % (step, loss, acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save and Load with TF Checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save weights and optimizer variables.\n",
    "# Create a dict of variables to save.\n",
    "vars_to_save = {\"W\": W, \"b\": b, \"optimizer\": optimizer}\n",
    "# TF Checkpoint, pass the dict as **kwargs.\n",
    "checkpoint = tf.train.Checkpoint(**vars_to_save)\n",
    "# TF CheckpointManager to manage saving parameters.\n",
    "saver = tf.train.CheckpointManager(\n",
    "      checkpoint, directory=\"./tf-example\", max_to_keep=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'./tf-example/ckpt-1'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Save variables.\n",
    "saver.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.09673191"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check weight value.\n",
    "np.mean(W.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset variables to test restore.\n",
    "W = tf.Variable(tf.random.normal([num_features, num_classes]), name=\"weight\")\n",
    "b = tf.Variable(tf.zeros([num_classes]), name=\"bias\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.0083419625"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check resetted weight value.\n",
    "np.mean(W.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x12cf965d0>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set checkpoint to load data.\n",
    "vars_to_load = {\"W\": W, \"b\": b, \"optimizer\": optimizer}\n",
    "checkpoint = tf.train.Checkpoint(**vars_to_load)\n",
    "# Restore variables from latest checkpoint.\n",
    "latest_ckpt = tf.train.latest_checkpoint(\"./tf-example\")\n",
    "checkpoint.restore(latest_ckpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.09673191"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Confirm that W has been correctly restored.\n",
    "np.mean(W.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2) TF Model\n",
    "\n",
    "Basic neural network with TF Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras import Model, layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MNIST dataset parameters.\n",
    "num_classes = 10 # 0 to 9 digits\n",
    "num_features = 784 # 28*28\n",
    "\n",
    "# Training parameters.\n",
    "learning_rate = 0.01\n",
    "training_steps = 1000\n",
    "batch_size = 256\n",
    "display_step = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create TF Model.\n",
    "class NeuralNet(Model):\n",
    "    # Set layers.\n",
    "    def __init__(self):\n",
    "        super(NeuralNet, self).__init__(name=\"NeuralNet\")\n",
    "        # First fully-connected hidden layer.\n",
    "        self.fc1 = layers.Dense(64, activation=tf.nn.relu)\n",
    "        # Second fully-connected hidden layer.\n",
    "        self.fc2 = layers.Dense(128, activation=tf.nn.relu)\n",
    "        # Third fully-connecter hidden layer.\n",
    "        self.out = layers.Dense(num_classes, activation=tf.nn.softmax)\n",
    "\n",
    "    # Set forward pass.\n",
    "    def __call__(self, x, is_training=False):\n",
    "        x = self.fc1(x)\n",
    "        x = self.out(x)\n",
    "        if not is_training:\n",
    "            # tf cross entropy expect logits without softmax, so only\n",
    "            # apply softmax when not training.\n",
    "            x = tf.nn.softmax(x)\n",
    "        return x\n",
    "\n",
    "# Build neural network model.\n",
    "neural_net = NeuralNet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cross-Entropy loss function.\n",
    "def cross_entropy(y_pred, y_true):\n",
    "    y_true = tf.cast(y_true, tf.int64)\n",
    "    crossentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_true, logits=y_pred)\n",
    "    return tf.reduce_mean(crossentropy)\n",
    "\n",
    "# Accuracy metric.\n",
    "def accuracy(y_pred, y_true):\n",
    "    # Predicted class is the index of highest score in prediction vector (i.e. argmax).\n",
    "    correct_prediction = tf.equal(tf.argmax(y_pred, 1), tf.cast(y_true, tf.int64))\n",
    "    return tf.reduce_mean(tf.cast(correct_prediction, tf.float32))\n",
    "\n",
    "# Adam optimizer.\n",
    "optimizer = tf.optimizers.Adam(learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimization process. \n",
    "def run_optimization(x, y):\n",
    "    # Wrap computation inside a GradientTape for automatic differentiation.\n",
    "    with tf.GradientTape() as g:\n",
    "        pred = neural_net(x, is_training=True)\n",
    "        loss = cross_entropy(pred, y)\n",
    "\n",
    "        # Compute gradients.\n",
    "        gradients = g.gradient(loss, neural_net.trainable_variables)\n",
    "\n",
    "        # Update W and b following gradients.\n",
    "        optimizer.apply_gradients(zip(gradients, neural_net.trainable_variables))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step: 100, loss: 2.188605, accuracy: 0.902344\n",
      "step: 200, loss: 2.182990, accuracy: 0.929688\n",
      "step: 300, loss: 2.180439, accuracy: 0.945312\n",
      "step: 400, loss: 2.178496, accuracy: 0.957031\n",
      "step: 500, loss: 2.177517, accuracy: 0.968750\n",
      "step: 600, loss: 2.177163, accuracy: 0.968750\n",
      "step: 700, loss: 2.177454, accuracy: 0.960938\n",
      "step: 800, loss: 2.177589, accuracy: 0.960938\n",
      "step: 900, loss: 2.176507, accuracy: 0.964844\n",
      "step: 1000, loss: 2.177557, accuracy: 0.960938\n"
     ]
    }
   ],
   "source": [
    "# Run training for the given number of steps.\n",
    "for step, (batch_x, batch_y) in enumerate(train_data.take(training_steps), 1):\n",
    "    # Run the optimization to update W and b values.\n",
    "    run_optimization(batch_x, batch_y)\n",
    "    \n",
    "    if step % display_step == 0:\n",
    "        pred = neural_net(batch_x, is_training=False)\n",
    "        loss = cross_entropy(pred, batch_y)\n",
    "        acc = accuracy(pred, batch_y)\n",
    "        print(\"step: %i, loss: %f, accuracy: %f\" % (step, loss, acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save and Load with TF Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save TF model.\n",
    "neural_net.save_weights(filepath=\"./tfmodel.ckpt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy: 0.101562\n"
     ]
    }
   ],
   "source": [
    "# Re-build neural network model with default values.\n",
    "neural_net = NeuralNet()\n",
    "# Test model performance.\n",
    "pred = neural_net(batch_x)\n",
    "print(\"accuracy: %f\" % accuracy(pred, batch_y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x13118b950>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load saved weights.\n",
    "neural_net.load_weights(filepath=\"./tfmodel.ckpt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy: 0.960938\n"
     ]
    }
   ],
   "source": [
    "# Test that weights loaded correctly.\n",
    "pred = neural_net(batch_x)\n",
    "print(\"accuracy: %f\" % accuracy(pred, batch_y))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
