Custom_Layer_with_Functional_API_saving_and_loading.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Custom Layer with Functional API: saving and loading.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyMK97s9F3QkYiSwCdDYuVIx",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jvishnuvardhan/Stackoverflow_Questions/blob/master/Custom_Layer_with_Functional_API_saving_and_loading.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RDTNssqfH-yk",
        "colab_type": "code",
        "outputId": "0a914339-8d5d-4ff3-9b93-149f62ce8cee",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 632
        }
      },
      "source": [
        "!pip install tensorflow==2.2.0rc4"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: tensorflow==2.2.0rc4 in /usr/local/lib/python3.6/dist-packages (2.2.0rc4)\n",
            "Requirement already satisfied: absl-py>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (0.9.0)\n",
            "Requirement already satisfied: protobuf>=3.8.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (3.10.0)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (3.2.1)\n",
            "Requirement already satisfied: grpcio>=1.8.6 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (1.28.1)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (1.12.0)\n",
            "Requirement already satisfied: keras-preprocessing>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (1.1.0)\n",
            "Requirement already satisfied: scipy==1.4.1; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (1.4.1)\n",
            "Requirement already satisfied: wheel>=0.26; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (0.34.2)\n",
            "Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (1.12.1)\n",
            "Requirement already satisfied: tensorboard<2.3.0,>=2.2.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (2.2.1)\n",
            "Requirement already satisfied: h5py<2.11.0,>=2.10.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (2.10.0)\n",
            "Requirement already satisfied: numpy<2.0,>=1.16.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (1.18.3)\n",
            "Requirement already satisfied: tensorflow-estimator<2.3.0,>=2.2.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (2.2.0)\n",
            "Requirement already satisfied: astunparse==1.6.3 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (1.6.3)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (1.1.0)\n",
            "Requirement already satisfied: gast==0.3.3 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (0.3.3)\n",
            "Requirement already satisfied: google-pasta>=0.1.8 in /usr/local/lib/python3.6/dist-packages (from tensorflow==2.2.0rc4) (0.2.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from protobuf>=3.8.0->tensorflow==2.2.0rc4) (46.1.3)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (3.2.1)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (1.6.0.post3)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (2.23.0)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (1.0.1)\n",
            "Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (1.7.2)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (0.4.1)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (2020.4.5.1)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (2.9)\n",
            "Requirement already satisfied: cachetools<3.2,>=2.0.0 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (3.1.1)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (0.2.8)\n",
            "Requirement already satisfied: rsa<4.1,>=3.1.4 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (4.0)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (1.3.0)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.6/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (0.4.8)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.6/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0rc4) (3.1.0)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YVX15QEzLNAg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "    def get_config(self):\n",
        "      config = super(CustomLayer, self).get_config()\n",
        "      config.update({'add_loss': self.add_loss})\n",
        "      return config"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DHjIV_JVH_2v",
        "colab_type": "code",
        "outputId": "c82049ca-48df-4a2d-c641-ae6ede6e7bde",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 714
        }
      },
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "class CustomLayer(tf.keras.layers.Layer):\n",
        "    \"\"\"Imaginary Layer that adds a custom loss in the call\"\"\"\n",
        "\n",
        "    def __init__(self, a):\n",
        "        super().__init__()\n",
        "        self.var = tf.Variable(a, name='var_a')\n",
        "\n",
        "    def call(self, inputs, training=False):\n",
        "      output = tf.reduce_sum(inputs * self.var, axis=-1)\n",
        "      self.add_loss(tf.reduce_mean(output))\n",
        "      return output\n",
        "    def get_config(self):\n",
        "      return {'a': self.var.numpy()}\n",
        "\n",
        "def get_model(input_dim: int) -> tf.keras.Model:\n",
        "    layer = CustomLayer(0.1)\n",
        "    inputs = tf.keras.Input((input_dim,), name=\"inputs\")\n",
        "    outputs = layer(inputs)\n",
        "    model = tf.keras.Model(inputs=inputs, outputs=outputs)\n",
        "    model.compile(optimizer='adam')\n",
        "    return model\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    num_data = 100\n",
        "    X, Y = np.random.randn(num_data, 1), np.random.randn(num_data, 1)\n",
        "\n",
        "    model = get_model(input_dim=X.shape[-1])\n",
        "    \n",
        "    model.summary()\n",
        "    print(\"model.losses\", model.losses)\n",
        "    model.fit(X,Y, epochs=2)\n",
        "    model.save('my_model',save_format='tf')\n",
        "    print('-'*65)\n",
        "    print('\\n-------------- After loading the model back -----------------\\n')\n",
        "    print('-'*65)\n",
        "    reconstructed_model = tf.keras.models.load_model('my_model',custom_objects={'CustomLayer':CustomLayer})  # ~~ breaks \n",
        "    reconstructed_model.summary()\n",
        "\n",
        "    # Let's check:\n",
        "    np.testing.assert_allclose(\n",
        "        model.predict(X),\n",
        "        reconstructed_model.predict(X)\n",
        "    )\n",
        "\n",
        "    print(\"model.losses:\", model.losses)\n",
        "    print(\"reconstructed_model.losses:\", reconstructed_model.losses)\n",
        "\n",
        "    reconstructed_model.fit(X,Y, epochs=2)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Model: \"model_29\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "inputs (InputLayer)          [(None, 1)]               0         \n",
            "_________________________________________________________________\n",
            "custom_layer_43 (CustomLayer (None,)                   1         \n",
            "=================================================================\n",
            "Total params: 1\n",
            "Trainable params: 1\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n",
            "model.losses [<tf.Tensor 'custom_layer_43/Mean:0' shape=() dtype=float32>]\n",
            "Epoch 1/2\n",
            "4/4 [==============================] - 0s 2ms/step - loss: 0.0055\n",
            "Epoch 2/2\n",
            "4/4 [==============================] - 0s 2ms/step - loss: 0.0038\n",
            "INFO:tensorflow:Assets written to: my_model/assets\n",
            "-----------------------------------------------------------------\n",
            "\n",
            "-------------- After loading the model back -----------------\n",
            "\n",
            "-----------------------------------------------------------------\n",
            "Model: \"model_29\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "inputs (InputLayer)          [(None, 1)]               0         \n",
            "_________________________________________________________________\n",
            "custom_layer_43 (CustomLayer (None,)                   1         \n",
            "=================================================================\n",
            "Total params: 1\n",
            "Trainable params: 1\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n",
            "model.losses: [<tf.Tensor 'model_29/custom_layer_43/Mean:0' shape=() dtype=float32>]\n",
            "reconstructed_model.losses: [<tf.Tensor 'model_29/custom_layer_43/Mean:0' shape=() dtype=float32>]\n",
            "Epoch 1/2\n",
            "4/4 [==============================] - 0s 2ms/step - loss: -0.0100\n",
            "Epoch 2/2\n",
            "4/4 [==============================] - 0s 1ms/step - loss: -0.0166\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kApoQpuEPB4z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}