tf_keras_for_Researchers_Crash_Course.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0IPokL4mivsz"
   },
   "outputs": [],
   "source": [
    "!pip install tensorflow==2.0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "rwREGJ7Wiyl9",
    "outputId": "fa8c0d6d-48b6-48dd-d82a-65a7d029089a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.0.0'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "tf.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xEuWqzjlPobA"
   },
   "source": [
    "# tf.keras for Researchers: Crash Course\n",
    "\n",
    "Are you a deep learning researcher? Wondering if all this TensorFlow 2.0 stuff is relevant to you?\n",
    "\n",
    "This notebook is a crash course on everything you need to know to use TensorFlow 2.0 for deep learning research."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cttOYWBtPx9f"
   },
   "source": [
    "1) The first class you need to know is `Layer`. A Layer encapsulates a state (weights) and some computation (defined in the `call` method).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "67-lEi5tjYwd"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras.layers import Layer\n",
    "\n",
    "\n",
    "class Linear(Layer):\n",
    "  \"\"\"y = w.x + b\"\"\"\n",
    "\n",
    "    def __init__(self, units=32, input_dim=32):\n",
    "        super(Linear, self).__init__()\n",
    "        w_init = tf.random_normal_initializer()\n",
    "        self.w = tf.Variable(\n",
    "          initial_value=w_init(shape=(input_dim, units), dtype='float32'),\n",
    "          trainable=True)\n",
    "        b_init = tf.zeros_initializer()\n",
    "        self.b = tf.Variable(\n",
    "          initial_value=b_init(shape=(units,), dtype='float32'),\n",
    "          trainable=True)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        return tf.matmul(inputs, self.w) + self.b\n",
    "\n",
    "\n",
    "# Instantiate our layer.\n",
    "linear_layer = Linear(4, 2)\n",
    "\n",
    "# The layer can be treated as a function.\n",
    "# Here we call it on some data.\n",
    "y = linear_layer(tf.ones((2, 2)))\n",
    "assert y.shape == (2, 4)\n",
    "\n",
    "# Weights are automatically tracked under the `weights` property.\n",
    "assert linear_layer.weights == [linear_layer.w, linear_layer.b]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<tf.Variable 'Variable:0' shape=(2, 4) dtype=float32, numpy=\n",
       " array([[ 0.01693215,  0.00674726,  0.06510697,  0.05872518],\n",
       "        [-0.04878018, -0.02374228, -0.05174647, -0.05652155]],\n",
       "       dtype=float32)>,\n",
       " <tf.Variable 'Variable:0' shape=(4,) dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "linear_layer.weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6O77qfbqPn06"
   },
   "source": [
    "2) The `add_weight` method gives you a shortcut for creating weights.\n",
    "\n",
    "3) It’s good practice to create weights in a separate `build` method, called lazily with the shape of the first inputs seen by your layer. Here, this pattern prevents us from having to specify `input_dim` in the constructor:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0SnzTy1Rjyi_"
   },
   "outputs": [],
   "source": [
    "class Linear(Layer):\n",
    "  \"\"\"y = w.x + b\"\"\"\n",
    "\n",
    "    def __init__(self, units=32):\n",
    "        super(Linear, self).__init__()\n",
    "        self.units = units\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        self.w = self.add_weight(shape=(input_shape[-1], self.units),\n",
    "                               initializer='random_normal',\n",
    "                               trainable=True)\n",
    "        self.b = self.add_weight(shape=(self.units,),\n",
    "                               initializer='random_normal',\n",
    "                               trainable=True)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        return tf.matmul(inputs, self.w) + self.b\n",
    "\n",
    "\n",
    "# Instantiate our lazy layer.\n",
    "linear_layer = Linear(4)\n",
    "\n",
    "# This will also call `build(input_shape)` and create the weights.\n",
    "y = linear_layer(tf.ones((2, 2)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MCyO_yLFQv4I"
   },
   "source": [
    "4) You can automatically retrieve the gradients of the weights of a layer by calling it inside a `GradientTape`.  Using these gradients, you can update the weights of the layer, either manually, or using an optimizer object. Of course, you can modify the gradients before using them, if you need to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 221
    },
    "colab_type": "code",
    "id": "UHOOlixcQ9Gl",
    "outputId": "8ba51007-d901-4223-a1b6-22e8ce7d6eea"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz\n",
      "11493376/11490434 [==============================] - 0s 0us/step\n",
      "0 2.418797731399536\n",
      "100 2.278667449951172\n",
      "200 2.156917095184326\n",
      "300 2.041379928588867\n",
      "400 2.018954277038574\n",
      "500 1.9337773323059082\n",
      "600 1.7980483770370483\n",
      "700 1.8048418760299683\n",
      "800 1.6862506866455078\n",
      "900 1.6376819610595703\n"
     ]
    }
   ],
   "source": [
    "# Prepare a dataset.\n",
    "(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()\n",
    "dataset = tf.data.Dataset.from_tensor_slices(\n",
    "    (x_train.reshape(60000, 784).astype('float32') / 255, y_train))\n",
    "dataset = dataset.shuffle(buffer_size=1024).batch(64)\n",
    "\n",
    "# Instantiate our linear layer (defined above) with 10 units.\n",
    "linear_layer = Linear(10)\n",
    "\n",
    "# Instantiate a logistic loss function that expects integer targets.\n",
    "loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "\n",
    "# Instantiate an optimizer.\n",
    "optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)\n",
    "\n",
    "# Iterate over the batches of the dataset.\n",
    "for step, (x, y) in enumerate(dataset):\n",
    "  \n",
    "  # Open a GradientTape.\n",
    "    with tf.GradientTape() as tape:\n",
    "\n",
    "        # Forward pass.\n",
    "        logits = linear_layer(x)\n",
    "\n",
    "        # Loss value for this batch.\n",
    "        loss = loss_fn(y, logits)\n",
    "     \n",
    "    # Get gradients of weights wrt the loss.\n",
    "    gradients = tape.gradient(loss, linear_layer.trainable_weights)\n",
    "\n",
    "    # Update the weights of our linear layer.\n",
    "    optimizer.apply_gradients(zip(gradients, linear_layer.trainable_weights))\n",
    "\n",
    "    # Logging.\n",
    "    if step % 100 == 0:\n",
    "    print(step, float(loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ghnbq7zeQJ6H"
   },
   "source": [
    "5) Weights created by layers can be either trainable or non-trainable. They're exposed in `trainable_weights` and `non_trainable_weights`. Here's a layer with a non-trainable weight:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "TdAF5xWFj9vp",
    "outputId": "b525c802-8f96-4086-9381-d3ac3ee54d11"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2. 2.]\n",
      "[4. 4.]\n"
     ]
    }
   ],
   "source": [
    "class ComputeSum(Layer):\n",
    "  \"\"\"Returns the sum of the inputs.\"\"\"\n",
    "\n",
    "    def __init__(self, input_dim):\n",
    "        super(ComputeSum, self).__init__()\n",
    "        # Create a non-trainable weight.\n",
    "        self.total = tf.Variable(initial_value=tf.zeros((input_dim,)),\n",
    "                               trainable=False)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        self.total.assign_add(tf.reduce_sum(inputs, axis=0))\n",
    "        return self.total  \n",
    "\n",
    "my_sum = ComputeSum(2)\n",
    "x = tf.ones((2, 2))\n",
    "\n",
    "y = my_sum(x)\n",
    "print(y.numpy())  # [2. 2.]\n",
    "\n",
    "y = my_sum(x)\n",
    "print(y.numpy())  # [4. 4.]\n",
    "\n",
    "assert my_sum.weights == [my_sum.total]\n",
    "assert my_sum.non_trainable_weights == [my_sum.total]\n",
    "assert my_sum.trainable_weights == []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2f-IqdBvQqVW"
   },
   "source": [
    "6) Layers can be recursively nested to create bigger computation blocks. Each layer will track the weights of its sublayers (both trainable and non-trainable)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4so7zVstj_Js"
   },
   "outputs": [],
   "source": [
    "# Let's reuse the Linear class\n",
    "# with a `build` method that we defined above.\n",
    "\n",
    "class MLP(Layer):\n",
    "    \"\"\"Simple stack of Linear layers.\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        super(MLP, self).__init__()\n",
    "        self.linear_1 = Linear(32)\n",
    "        self.linear_2 = Linear(32)\n",
    "        self.linear_3 = Linear(10)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        x = self.linear_1(inputs)\n",
    "        x = tf.nn.relu(x)\n",
    "        x = self.linear_2(x)\n",
    "        x = tf.nn.relu(x)\n",
    "        return self.linear_3(x)\n",
    "\n",
    "mlp = MLP()\n",
    "\n",
    "# The first call to the `mlp` object will create the weights.\n",
    "y = mlp(tf.ones(shape=(3, 64)))\n",
    "\n",
    "# Weights are recursively tracked.\n",
    "assert len(mlp.weights) == 6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Eg-3_LxhUhGX"
   },
   "source": [
    "7) Layers can create losses during the forward pass. This is especially useful for regularization losses. The losses created by sublayers are recursively tracked by the parent layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "cnGNfHAWTvGe",
    "outputId": "2b2aadc5-60a9-4c8a-c2a9-db42f38890d5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<tf.Tensor: id=62181, shape=(), dtype=float32, numpy=0.17264488>]\n"
     ]
    }
   ],
   "source": [
    "class ActivityRegularization(Layer):\n",
    "    \"\"\"Layer that creates an activity sparsity regularization loss.\"\"\"\n",
    "  \n",
    "    def __init__(self, rate=1e-2):\n",
    "        super(ActivityRegularization, self).__init__()\n",
    "        self.rate = rate\n",
    "\n",
    "    def call(self, inputs):\n",
    "        # We use `add_loss` to create a regularization loss\n",
    "        # that depends on the inputs.\n",
    "        self.add_loss(self.rate * tf.reduce_sum(inputs))\n",
    "        return inputs\n",
    "\n",
    "\n",
    "# Let's use the loss layer in a MLP block.\n",
    "\n",
    "class SparseMLP(Layer):\n",
    "    \"\"\"Stack of Linear layers with a sparsity regularization loss.\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        super(SparseMLP, self).__init__()\n",
    "        self.linear_1 = Linear(32)\n",
    "        self.regularization = ActivityRegularization(1e-2)\n",
    "        self.linear_3 = Linear(10)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        x = self.linear_1(inputs)\n",
    "        x = tf.nn.relu(x)\n",
    "        x = self.regularization(x)\n",
    "        return self.linear_3(x)\n",
    "\n",
    "\n",
    "mlp = SparseMLP()\n",
    "y = mlp(tf.ones((10, 10)))\n",
    "\n",
    "print(mlp.losses)  # List containing one float32 scalar"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "p69GSoulV632"
   },
   "source": [
    "8) These losses are cleared by the top-level layer at the start of each forward pass -- they don't accumulate. `layer.losses` always contain only the losses created during the *last* forward pass. You would typically use these losses by summing them before computing your gradients when writing a training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 187
    },
    "colab_type": "code",
    "id": "WH8alljNVj2J",
    "outputId": "ad1d50bb-3260-4164-d353-859c2c6e4a66"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 5.764951705932617\n",
      "100 2.5185461044311523\n",
      "200 2.4115047454833984\n",
      "300 2.338585138320923\n",
      "400 2.3432374000549316\n",
      "500 2.3476529121398926\n",
      "600 2.3310372829437256\n",
      "700 2.3395984172821045\n",
      "800 2.3197717666625977\n",
      "900 2.3151700496673584\n"
     ]
    }
   ],
   "source": [
    "# Losses correspond to the *last* forward pass.\n",
    "mlp = SparseMLP()\n",
    "mlp(tf.ones((10, 10)))\n",
    "assert len(mlp.losses) == 1\n",
    "mlp(tf.ones((10, 10)))\n",
    "assert len(mlp.losses) == 1  # No accumulation.\n",
    "\n",
    "# Let's demonstrate how to use these losses in a training loop.\n",
    "\n",
    "# Prepare a dataset.\n",
    "(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()\n",
    "dataset = tf.data.Dataset.from_tensor_slices(\n",
    "    (x_train.reshape(60000, 784).astype('float32') / 255, y_train))\n",
    "dataset = dataset.shuffle(buffer_size=1024).batch(64)\n",
    "\n",
    "# A new MLP.\n",
    "mlp = SparseMLP()\n",
    "\n",
    "# Loss and optimizer.\n",
    "loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)\n",
    "\n",
    "for step, (x, y) in enumerate(dataset):\n",
    "    with tf.GradientTape() as tape:\n",
    "\n",
    "        # Forward pass.\n",
    "        logits = mlp(x)\n",
    "\n",
    "        # External loss value for this batch.\n",
    "        loss = loss_fn(y, logits)\n",
    "\n",
    "        # Add the losses created during the forward pass.\n",
    "        loss += sum(mlp.losses)\n",
    "\n",
    "        # Get gradients of weights wrt the loss.\n",
    "        gradients = tape.gradient(loss, mlp.trainable_weights)\n",
    "\n",
    "    # Update the weights of our linear layer.\n",
    "    optimizer.apply_gradients(zip(gradients, mlp.trainable_weights))\n",
    "\n",
    "    # Logging.\n",
    "    if step % 100 == 0:\n",
    "    print(step, float(loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jcL_zMFtZ1R8"
   },
   "source": [
    "9) Running eagerly is great for debugging, but you will get better performance by compiling your computation into static graphs. Static graphs are a researcher's best friends. You can compile any function by wrapping it in a `tf.function` decorator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 187
    },
    "colab_type": "code",
    "id": "2XmEPQhhY4ef",
    "outputId": "1dc339a0-ad37-4537-af10-b8374cf433e0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 2.3017988204956055\n",
      "100 2.30670166015625\n",
      "200 2.3061323165893555\n",
      "300 2.3000917434692383\n",
      "400 2.298370838165283\n",
      "500 2.298713207244873\n",
      "600 2.303539276123047\n",
      "700 2.3063406944274902\n",
      "800 2.293882369995117\n",
      "900 2.299696922302246\n"
     ]
    }
   ],
   "source": [
    "# Prepare our layer, loss, and optimizer.\n",
    "mlp = MLP()\n",
    "loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)\n",
    "\n",
    "# Create a training step function.\n",
    "\n",
    "@tf.function  # Make it fast.\n",
    "def train_on_batch(x, y):\n",
    "    with tf.GradientTape() as tape:\n",
    "        logits = mlp(x)\n",
    "        loss = loss_fn(y, logits)\n",
    "        gradients = tape.gradient(loss, mlp.trainable_weights)\n",
    "    optimizer.apply_gradients(zip(gradients, mlp.trainable_weights))\n",
    "    return loss\n",
    "\n",
    "# Prepare a dataset.\n",
    "(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()\n",
    "dataset = tf.data.Dataset.from_tensor_slices(\n",
    "    (x_train.reshape(60000, 784).astype('float32') / 255, y_train))\n",
    "dataset = dataset.shuffle(buffer_size=1024).batch(64)\n",
    "\n",
    "for step, (x, y) in enumerate(dataset):\n",
    "    loss = train_on_batch(x, y)\n",
    "    if step % 100 == 0:\n",
    "        print(step, float(loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8dNI9xczb8ej"
   },
   "source": [
    "10) Some layers, in particular the `BatchNormalization` layer and the `Dropout` layer, have different behaviors during training and inference. For such layers, it is standard practice to expose a `training` (boolean) argument in the `call` method.\n",
    "\n",
    "By exposing this argument in `call`, you enable the built-in training and evaluation loops (e.g. `fit`) to correctly use the layer in training and inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uwttOFombD1F"
   },
   "outputs": [],
   "source": [
    "class Dropout(Layer):\n",
    "  \n",
    "  def __init__(self, rate):\n",
    "    super(Dropout, self).__init__()\n",
    "    self.rate = rate\n",
    "\n",
    "  def call(self, inputs, training=None):\n",
    "    if training:\n",
    "      return tf.nn.dropout(inputs, rate=self.rate)\n",
    "    return inputs\n",
    "  \n",
    "\n",
    "class MLPWithDropout(Layer):\n",
    "\n",
    "  def __init__(self):\n",
    "      super(MLPWithDropout, self).__init__()\n",
    "      self.linear_1 = Linear(32)\n",
    "      self.dropout = Dropout(0.5)\n",
    "      self.linear_3 = Linear(10)\n",
    "\n",
    "  def call(self, inputs, training=None):\n",
    "      x = self.linear_1(inputs)\n",
    "      x = tf.nn.relu(x)\n",
    "      x = self.dropout(x, training=training)\n",
    "      return self.linear_3(x)\n",
    "    \n",
    "mlp = MLPWithDropout()\n",
    "y_train = mlp(tf.ones((2, 2)), training=True)\n",
    "y_test = mlp(tf.ones((2, 2)), training=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "THRG48kgmQh1"
   },
   "source": [
    "11) You have many built-in layers available, from `Dense` to `Conv2D` to `LSTM` to fancier ones like `Conv2DTranspose` or `ConvLSTM2D`. Be smart about reusing built-in functionality."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "euLbXDzgnFti"
   },
   "source": [
    "12) To build deep learning models, you don't have to use object-oriented programming all the time. All layers we've seen so far can also be composed functionally, like this (we call it the \"Functional API\"):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Wb9ZT-eBoBPe"
   },
   "outputs": [],
   "source": [
    "# We use an `Input` object to describe the shape and dtype of the inputs.\n",
    "# This is the deep learning equivalent of *declaring a type*.\n",
    "# The shape argument is per-sample; it does not include the batch size.\n",
    "# The functional API focused on defining per-sample transformations.\n",
    "# The model we create will automatically batch the per-sample transformations,\n",
    "# so that it can be called on batches of data.\n",
    "inputs = tf.keras.Input(shape=(16,))\n",
    "\n",
    "# We call layers on these \"type\" objects\n",
    "# and they return updated types (new shapes/dtypes).\n",
    "x = Linear(32)(inputs) # We are reusing the Linear layer we defined earlier.\n",
    "x = Dropout(0.5)(x) # We are reusing the Dropout layer we defined earlier.\n",
    "outputs = Linear(10)(x)\n",
    "\n",
    "# A functional `Model` can be defined by specifying inputs and outputs.\n",
    "# A model is itself a layer like any other.\n",
    "model = tf.keras.Model(inputs, outputs)\n",
    "\n",
    "# A functional model already has weights, before being called on any data.\n",
    "# That's because we defined its input shape in advance (in `Input`).\n",
    "assert len(model.weights) == 4\n",
    "\n",
    "# Let's call our model on some data, for fun.\n",
    "y = model(tf.ones((2, 16)))\n",
    "assert y.shape == (2, 10)\n",
    "\n",
    "# You can pass a `training` argument in `__call__`\n",
    "# (it will get passed down to the Dropout layer).\n",
    "y = model(tf.ones((2, 16)), training=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ABcc4nAipMyp"
   },
   "source": [
    "The Functional API tends to be more concise than subclassing, and provides a few other advantages (generally the same advantages that functional, typed languages provide over untyped OO development). However, it can only be used to define DAGs of layers -- recursive networks should be defined as `Layer` subclasses instead.\n",
    "\n",
    "Learn more about the Functional API [here](https://www.tensorflow.org/alpha/guide/keras/functional).\n",
    "\n",
    "In your research workflows, you may often find yourself mix-and-matching OO models and Functional models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XhKL2JsRqHq3"
   },
   "source": [
    "### That's all you need to get started with reimplementing most deep learning research papers in TensorFlow 2.0 and Keras!\n",
    "\n",
    "## Now for a quick research example: hypernetworks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lFnm79N-qwP3"
   },
   "source": [
    "Let's put these concepts into practice with a simple end-to-end example.\n",
    "\n",
    "A [hypernetwork](http://blog.otoro.net/2016/09/28/hyper-networks/) is a deep neural network whose weights are generated by another network (usually smaller).\n",
    "\n",
    "Let's implement a really trivial hypernetwork: we'll take the `Linear` layer we defined earlier, and we'll use it to generate the weights of... another `Linear` layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "3PlPkpKnr14L",
    "outputId": "f3485c73-8b7d-4202-aea8-10f186451fb2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 1.7940431833267212\n",
      "100 4.222971867615973\n",
      "200 3.1491951361929185\n",
      "300 2.6909681088329234\n",
      "400 2.355490146474159\n",
      "500 2.1030502786059238\n",
      "600 1.9476623649384823\n",
      "700 1.8217411460702804\n",
      "800 1.6731741473888775\n",
      "900 1.5703826624686852\n",
      "1000 1.5540820519197744\n"
     ]
    }
   ],
   "source": [
    "input_dim = 784\n",
    "classes = 10\n",
    "\n",
    "# The model we'll actually use (the hypernetwork).\n",
    "outer_model = Linear(classes)\n",
    "\n",
    "# It doesn't need to create its own weights, so let's mark it as already built.\n",
    "# That way, calling `outer_model` won't create new variables.\n",
    "outer_model.built = True\n",
    "\n",
    "# The model that generates the weights of the model above.\n",
    "inner_model = Linear(input_dim * classes + classes)\n",
    "\n",
    "# Loss and optimizer.\n",
    "loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)\n",
    "\n",
    "# Prepare a dataset.\n",
    "(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()\n",
    "dataset = tf.data.Dataset.from_tensor_slices(\n",
    "    (x_train.reshape(60000, 784).astype('float32') / 255, y_train))\n",
    "\n",
    "# We'll use a batch size of 1 for this experiment.\n",
    "dataset = dataset.shuffle(buffer_size=1024).batch(1)\n",
    "  \n",
    "losses = []  # Keep track of the losses over time.\n",
    "for step, (x, y) in enumerate(dataset):\n",
    "  with tf.GradientTape() as tape:\n",
    "\n",
    "    # Predict weights for the outer model.\n",
    "    weights_pred = inner_model(x)\n",
    "\n",
    "    # Reshape them to the expected shapes for w and b for the outer model.\n",
    "    w_pred = tf.reshape(weights_pred[:, :-classes], (input_dim, classes))\n",
    "    b_pred = tf.reshape(weights_pred[:, -classes:], (classes,))\n",
    "    \n",
    "    # Set the weight predictions as the weight variables on the outer model.\n",
    "    outer_model.w = w_pred\n",
    "    outer_model.b = b_pred\n",
    "    \n",
    "    # Inference on the outer model.\n",
    "    preds = outer_model(x)\n",
    "    loss = loss_fn(y, preds)\n",
    "\n",
    "  # Train only inner model.\n",
    "  grads = tape.gradient(loss, inner_model.trainable_weights)\n",
    "  optimizer.apply_gradients(zip(grads, inner_model.trainable_weights))\n",
    "  \n",
    "  # Logging.\n",
    "  losses.append(float(loss))\n",
    "  if step % 100 == 0:\n",
    "    print(step, sum(losses) / len(losses))\n",
    "    \n",
    "  # Stop after 1000 steps.\n",
    "  if step >= 1000:\n",
    "    break"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "tf.keras for Researchers: Crash Course.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "venv3.6_tf2",
   "language": "python",
   "name": "venv3.6_tf2"
  },
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}