Loading_subclassed_model_in_tf_format.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "38251 tf-nightly.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "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/Loading_subclassed_model_in_tf_format.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2x3ggVPtB0NC",
        "colab_type": "code",
        "outputId": "949d1dc3-05b1-413e-e344-e87a364c136c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 768
        }
      },
      "source": [
        "pip install tf-nightly"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting tf-nightly\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/19/71/67a8c442164b92e77685438f14483cb5a1e91bcaccf0dd1da2a23647c65e/tf_nightly-2.2.0.dev20200406-cp36-cp36m-manylinux2010_x86_64.whl (517.1MB)\n",
            "\u001b[K     |████████████████████████████████| 517.2MB 28kB/s \n",
            "\u001b[?25hRequirement already satisfied: keras-preprocessing>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.1.0)\n",
            "Requirement already satisfied: scipy==1.4.1; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.4.1)\n",
            "Requirement already satisfied: google-pasta>=0.1.8 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (0.2.0)\n",
            "Collecting tb-nightly<2.4.0a0,>=2.3.0a0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/a6/63/a0b69e4c814581e52bfa76ec456fe7b910242b9497d8d57732ac1729da53/tb_nightly-2.3.0a20200406-py3-none-any.whl (2.9MB)\n",
            "\u001b[K     |████████████████████████████████| 2.9MB 37.7MB/s \n",
            "\u001b[?25hRequirement already satisfied: h5py<2.11.0,>=2.10.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (2.10.0)\n",
            "Requirement already satisfied: numpy<2.0,>=1.16.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.18.2)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (3.2.0)\n",
            "Collecting tf-estimator-nightly\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/51/09/88a430a478bc75d8cdc46edd2b3b1151380c8b7e72bf6e51080f88225109/tf_estimator_nightly-2.3.0.dev2020040601-py2.py3-none-any.whl (455kB)\n",
            "\u001b[K     |████████████████████████████████| 460kB 47.9MB/s \n",
            "\u001b[?25hRequirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.12.0)\n",
            "Requirement already satisfied: wheel>=0.26; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (0.34.2)\n",
            "Requirement already satisfied: protobuf>=3.9.2 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (3.10.0)\n",
            "Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.12.1)\n",
            "Requirement already satisfied: astunparse==1.6.3 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.6.3)\n",
            "Requirement already satisfied: gast==0.3.3 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (0.3.3)\n",
            "Requirement already satisfied: grpcio>=1.8.6 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.27.2)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (1.1.0)\n",
            "Requirement already satisfied: absl-py>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from tf-nightly) (0.9.0)\n",
            "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (46.1.3)\n",
            "Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (1.7.2)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (0.4.1)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (1.6.0.post2)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (1.0.1)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (3.2.1)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.6/dist-packages (from tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (2.21.0)\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->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (3.1.1)\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->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (4.0)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (0.2.8)\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->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (1.3.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (2019.11.28)\n",
            "Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (3.0.4)\n",
            "Requirement already satisfied: idna<2.9,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (2.8)\n",
            "Requirement already satisfied: urllib3<1.25,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (1.24.3)\n",
            "Requirement already satisfied: pyasn1>=0.1.3 in /usr/local/lib/python3.6/dist-packages (from rsa<4.1,>=3.1.4->google-auth<2,>=1.6.3->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (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->tb-nightly<2.4.0a0,>=2.3.0a0->tf-nightly) (3.1.0)\n",
            "Installing collected packages: tb-nightly, tf-estimator-nightly, tf-nightly\n",
            "Successfully installed tb-nightly-2.3.0a20200406 tf-estimator-nightly-2.3.0.dev2020040601 tf-nightly-2.2.0.dev20200406\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1iQmirddB0w5",
        "colab_type": "code",
        "outputId": "f2863eae-1965-4ea1-d61f-2833c23b88a9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "import tensorflow as tf\n",
        "tf.__version__"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'2.2.0-dev20200406'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3VBuEPopKz55",
        "colab_type": "code",
        "outputId": "25f6fa2c-0ca1-4d72-b3b8-06dbc180a298",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from sklearn.datasets import fetch_california_housing\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "housing = fetch_california_housing()\n",
        "X_train_full, X_test, y_train_full, y_test = train_test_split(housing.data, housing.target, random_state=42)\n",
        "X_train, X_valid, y_train, y_valid = train_test_split(X_train_full, y_train_full, random_state=42)\n",
        "X_train_A, X_train_B = X_train[:, :5], X_train[:, 2:]"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading Cal. housing from https://ndownloader.figshare.com/files/5976036 to /root/scikit_learn_data\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NIBB201sK4Yl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class WideAndDeepModel(keras.models.Model):\n",
        "    def __init__(self, units=30, activation=\"relu\", **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "        self.hidden1 = keras.layers.Dense(units, activation=activation, name='Dense-1')\n",
        "        self.hidden2 = keras.layers.Dense(units, activation=activation, name='Dense-2')\n",
        "        self.main_output = keras.layers.Dense(1, name='main_output')\n",
        "        self.aux_output = keras.layers.Dense(1, name='aux_output')\n",
        "        \n",
        "    def call(self, inputs):\n",
        "        input_A, input_B = inputs\n",
        "        hidden1 = self.hidden1(input_B)\n",
        "        hidden2 = self.hidden2(hidden1)\n",
        "        concat = keras.layers.concatenate([input_A, hidden2])\n",
        "        main_output = self.main_output(concat)\n",
        "        aux_output = self.aux_output(hidden2)\n",
        "        return main_output, aux_output\n",
        "\n",
        "model = WideAndDeepModel(30, activation=\"relu\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Sr41hemSK6MD",
        "colab_type": "code",
        "outputId": "dbfa1e9a-e68a-430d-86e5-9530a00274df",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "model.compile(loss=\"mse\", loss_weights=[0.9, 0.1], optimizer=keras.optimizers.SGD(lr=1e-3))\n",
        "history = model.fit((X_train_A, X_train_B), (y_train, y_train), epochs=1)\n",
        "model.save_weights('model_weights', save_format='tf')\n",
        "model.save_weights('model_weights_h5.h5', save_format='h5')"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "363/363 [==============================] - 0s 1ms/step - loss: nan - output_1_loss: nan - output_2_loss: nan\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jlG93uC0K71g",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "new_model = WideAndDeepModel(30, activation=\"relu\")\n",
        "new_model.compile(loss=\"mse\", loss_weights=[0.9, 0.1], optimizer=keras.optimizers.SGD(lr=1e-3))\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qIT0Sv2T6FjL",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "19dedcf1-8c3b-4411-db7e-0e2d98b27f64"
      },
      "source": [
        "# This initializes the variables used by the optimizers,\n",
        "# as well as any stateful metric variables\n",
        "new_model.train_on_batch((X_train_A[:1], X_train_B[:1]),((y_train[:1], y_train[:1])))\n",
        "\n",
        "# Load the state of the old model (to load weights for all layers)\n",
        "new_model.load_weights('model_weights')"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7fa505736a90>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hoQpnwTrLBHj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "new_model_h5 = WideAndDeepModel(30, activation=\"relu\")\n",
        "new_model_h5.compile(loss=\"mse\", loss_weights=[0.9, 0.1], optimizer=keras.optimizers.SGD(lr=1e-3))\n",
        "\n",
        "new_model_h5.train_on_batch((X_train_A[:1], X_train_B[:1]),((y_train[:1], y_train[:1])))\n",
        "# Load the state of the old model (to load weights for all layers)\n",
        "new_model_h5.load_weights('model_weights_h5.h5',by_name = True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t9PLQrbqANf1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "TL/DR: So if you want to load a model saved in *.tf format, then select `by_name=False`. If you want to load a model saved in *.h5 format, then select `by_name=True`. Please check the [gist here](https://colab.research.google.com/gist/jvishnuvardhan/c0c9808471249fb63298d9273b827f1e/38251-tf-nightly.ipynb).\n",
        "\n",
        "There is small difference between loading a saved model when the model was saved in *.tf format (tensorflow format) or in *.h5 format. The following error clearly shows what we need to do when we load a tensorflow-formatted weights (it need by_name =False). \n",
        "\n",
        "`NotImplementedError: Weights may only be loaded based on topology into Models when loading TensorFlow-formatted weights (got by_name=True to load_weights).`\n",
        "\n",
        "Check the description of `load_weights` on [TF website](https://www.tensorflow.org/api_docs/python/tf/keras/Model?version=nightly#load_weights). \n",
        "\n",
        "`Only topological loading (by_name=False) is supported when loading weights from the TensorFlow format. Note that topological loading differs slightly between TensorFlow and HDF5 formats for user-defined classes inheriting from tf.keras.Model: HDF5 loads based on a flattened list of weights, while the TensorFlow format loads based on the object-local names of attributes to which layers are assigned in the Model's constructor. `"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}