MAML_Implementation_FINAL_CODE.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "MAML_Implementation_FINAL_CODE.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyOSAgn+J9JeYy0SWkjTtwsx",
      "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/alorber/MAML-Implementation/blob/main/MAML_Implementation_FINAL_CODE.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k5yhb0E6kfvI"
      },
      "source": [
        "# Deep Learning Midterm Project\n",
        "## Implementation of [*Model-Agnostic Meta Learning for Fast Adaptation of Deep Network* by Finn](https://arxiv.org/abs/1703.03400)\n",
        "### Andrew Lorber and Mark Koszykowski"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Uv4Dvt5WkTer"
      },
      "source": [
        "# Imports\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from time import process_time\n",
        "from google.colab import files"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ESUO8aRdkpIB"
      },
      "source": [
        "## Sinusoid class\n",
        "Amplitude varies within [0.1, 5.0]\n",
        "\n",
        "Phase varies within [0, π]"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tTHO9u2PkqE1"
      },
      "source": [
        "# Creates a sinusoid from the given distribution\n",
        "class SineWave:\n",
        "    def __init__(self):\n",
        "        self.amplitude = np.random.uniform(low=0.1, high=5.0)\n",
        "        self.phase = np.random.uniform(0, np.pi)\n",
        "        \n",
        "    # returns k samples from the sine wave\n",
        "    # x is sampled uniformly from [-5.0, 5.0]\n",
        "    def sample(self, k):\n",
        "        x = tf.random.uniform(shape=(k,1), minval=-5.0, maxval=5.0)\n",
        "        y = self.amplitude * (tf.sin(x - self.phase))\n",
        "        return x,y\n",
        "\n",
        "    # Returns 1000 points of sinewave for graph\n",
        "    def truth(self):\n",
        "        x = tf.linspace(-5.0, 5.0, 1000)\n",
        "        y = self.amplitude * (tf.sin(x - self.phase))\n",
        "        return tf.reshape(x, (1000,1)),y"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ikaQbG03kuGC"
      },
      "source": [
        "## MAML Class\n",
        "*Subclass of the Keras Model Class*\n",
        "\n",
        "---\n",
        "\n",
        "The model consists of two hidden layers of size 40 with ReLU nonlinearities.\n",
        "\n",
        "The class has functions to feed forward input, meta train with MAML, and fine-tune to a specific sinusoid."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OFMQPiKxkwZn"
      },
      "source": [
        "# MAML Class\n",
        "class MAML(keras.Model):\n",
        "    # Builds the model as defined in the paper: 2 hidden layers of size 40 with ReLU\n",
        "    def __init__(self, k=10, alpha=0.01, beta=0.01):\n",
        "        super().__init__()\n",
        "\n",
        "        # Defines layers of model\n",
        "        self.hidden_layer_1 = keras.layers.Dense(40, input_shape=(1,))\n",
        "        self.hidden_layer_2 = keras.layers.Dense(40)\n",
        "        self.output_layer = keras.layers.Dense(1)\n",
        "\n",
        "        # Sets optimizer and loss function\n",
        "        self.optimizer = keras.optimizers.Adam(learning_rate=beta)\n",
        "        self.loss_fnc = keras.losses.mean_squared_error\n",
        "        \n",
        "        self.k = k\n",
        "        self.alpha = alpha\n",
        "\n",
        "    # Passes input, x, through the neural network and returns output\n",
        "    def f(self, x):\n",
        "        x_h_2 = keras.activations.relu(self.hidden_layer_1(x))\n",
        "        x_out = keras.activations.relu(self.hidden_layer_2(x_h_2))\n",
        "        y = self.output_layer(x_out)\n",
        "        return y\n",
        "\n",
        "    # Trains model following Algorithm 2 of paper\n",
        "    def meta_train(self, epochs=10000, tasks = 25):\n",
        "        num_epochs = epochs \n",
        "        num_tasks = tasks   # Number of tasks per epoch\n",
        "        total_loss = []     # Log of losses for graphing\n",
        "\n",
        "        # Step 1) Randomly initialize theta\n",
        "        # Completed during initial forward pass\n",
        "\n",
        "        # Calculates time per 100 epochs\n",
        "        t1 = process_time()\n",
        "\n",
        "        # Step 2) For each epoch do:\n",
        "        for epoch in range(num_epochs):\n",
        "            # Track progress\n",
        "            if epoch % 100 == 0 and epoch != 0:\n",
        "                t2 = process_time()\n",
        "                print(f\"Meta Learning Epoch: {epoch} - Time Per 100 Epochs: {t2 - t1} seconds\")\n",
        "                t1 = process_time()\n",
        "\n",
        "            losses = []     # Contains losses of each task for this epoch\n",
        "\n",
        "            # Step 3) Sample batch of tasks T_i ~ p(T)\n",
        "            sine_tasks = [SineWave() for _ in range(num_tasks)]\n",
        "\n",
        "            # Step 4) For each T_i do:\n",
        "            with tf.GradientTape() as outerTape:\n",
        "                for sinusoid in sine_tasks:\n",
        "                    # Step 5) Sample K datapoints D from T_i\n",
        "                    x_train, y_train = sinusoid.sample(self.k)\n",
        "\n",
        "                    # Step 6) Evaluate gradient w.r.t theta of Loss(f_theta)\n",
        "                    with tf.GradientTape() as innerTape:\n",
        "                        y_hat = self.f(x_train)\n",
        "                        loss = tf.reduce_mean(self.loss_fnc(y_train, y_hat))\n",
        "                    gradient = innerTape.gradient(loss, self.trainable_variables)\n",
        "\n",
        "                    # Step 7) Compute adapted parameters thetas_prime\n",
        "                    # Want to do this on copy of the model, so original thetas don't change\n",
        "                    model_copy = MAML()\n",
        "                    # Can't \"set_weights\" until weights are initialized (happens on first forward pass)\n",
        "                    model_copy.f(x_train)   \n",
        "                    model_copy.set_weights(self.get_weights())\n",
        "                    \n",
        "                    # Tensorflow does not allow nested \"apply_gradients\", \n",
        "                    #   so we used the following link for help:\n",
        "                    # https://stackoverflow.com/questions/58856784/nested-gradient-tape-in-function-tf2-0\n",
        "                    \n",
        "                    grad = 0\n",
        "                    for layer in range(len(model_copy.layers)):\n",
        "                        # Calculates adapted parameters w/ gradient descent\n",
        "                        # theta' = theta - alpha * gradients \n",
        "                        model_copy.layers[layer].kernel = tf.subtract(self.layers[layer].kernel, tf.Variable(tf.multiply(self.alpha, gradient[grad])))\n",
        "                        model_copy.layers[layer].bias = tf.subtract(self.layers[layer].bias, tf.Variable(tf.multiply(self.alpha, gradient[grad+1])))\n",
        "\n",
        "                        grad += 2                        \n",
        "\n",
        "                    # Step 8) Sample datapoints D' from T_i for the meta update\n",
        "                    x_train, y_train = sinusoid.sample(self.k)\n",
        "                    y_hat = model_copy.f(x_train)\n",
        "                    loss = tf.reduce_mean(model_copy.loss_fnc(y_train, y_hat))\n",
        "                    losses.append(loss)\n",
        "                \n",
        "                # Step 9) End for\n",
        "                loss = tf.reduce_sum(losses)\n",
        "                total_loss.append(loss)\n",
        "            \n",
        "            gradient = outerTape.gradient(loss, self.trainable_variables)\n",
        "            \n",
        "            # Step 10) Update thetas using D'_i and Loss(T_i)\n",
        "            self.optimizer.apply_gradients(zip(gradient, self.trainable_variables))\n",
        "            \n",
        "        # Step 11) End for\n",
        "        t2 = process_time()\n",
        "        print(f\"Meta Learning Epoch: {epochs} - Time Per 100 Epochs: {t2 - t1} seconds\")\n",
        "        \n",
        "        # Plots loss over epochs\n",
        "        fig = plt.figure()\n",
        "        plt.plot(range(epochs), total_loss)\n",
        "        plt.title(\"Meta-Loss per Epoch\")\n",
        "        fig.show()\n",
        "        fig.savefig(\"Loss.pdf\", bbox_inches='tight', pad_inches=0.25)\n",
        "        files.download(\"Loss.pdf\")\n",
        "\n",
        "    \n",
        "    # Fine-tunes model (and baseline if given) on randomly generated sine-wave\n",
        "    def fineTune(self, k, grad_steps=10, baseline = False):\n",
        "        # Creates new sine-wave & gets data-sets\n",
        "        wave = SineWave()\n",
        "        x_test, y_test = wave.sample(k)\n",
        "        x_truth, y_truth = wave.truth()\n",
        "\n",
        "        # MAML\n",
        "        # -----\n",
        "\n",
        "        # Creates copy of model, so it can be fine-tuned to different sine-waves for testing\n",
        "        maml_copy = MAML()\n",
        "        maml_copy.f(x_test)   \n",
        "        maml_copy.set_weights(self.get_weights())\n",
        "\n",
        "        # Saves the regression line pre-fine-tuning\n",
        "        y_preupdate_maml = maml_copy.f(x_truth)\n",
        "\n",
        "        # Applies \"grad_steps\" gradient steps\n",
        "        for i in range(grad_steps):\n",
        "            # Calculates gradient\n",
        "            with tf.GradientTape() as tape:\n",
        "                y_hat = maml_copy.f(x_test)\n",
        "                loss = tf.reduce_mean(maml_copy.loss_fnc(y_test, y_hat))\n",
        "            gradient = tape.gradient(loss, maml_copy.trainable_variables)\n",
        "            maml_copy.optimizer.apply_gradients(zip(gradient, maml_copy.trainable_variables))\n",
        "\n",
        "            # Saves regression line after 1st gradient step\n",
        "            if(i == 0):\n",
        "                y_1grad_maml = maml_copy.f(x_truth)\n",
        "        \n",
        "        # Saves regression line after \"grad_steps\" gradient steps\n",
        "        y_ngrad_maml = maml_copy.f(x_truth)\n",
        "\n",
        "        # Baseline\n",
        "        # ---------\n",
        "\n",
        "        if(baseline):\n",
        "            # Creates copy of model, so it can be fine-tuned to different sine-waves for testing\n",
        "            baseline_copy = Baseline()\n",
        "            baseline_copy.f(x_test)   \n",
        "            baseline_copy.set_weights(baseline.get_weights())\n",
        "\n",
        "            # Saves the regression line pre-fine-tuning\n",
        "            y_preupdate_baseline = baseline_copy.f(x_truth)\n",
        "\n",
        "            # Applies \"grad_steps\" gradient steps\n",
        "            for i in range(grad_steps):\n",
        "                # Calculates gradient\n",
        "                with tf.GradientTape() as tape:\n",
        "                    y_hat = baseline_copy.f(x_test)\n",
        "                    loss = tf.reduce_mean(baseline_copy.loss_fnc(y_test, y_hat))\n",
        "                gradient = tape.gradient(loss, baseline_copy.trainable_variables)\n",
        "                baseline_copy.optimizer.apply_gradients(zip(gradient, baseline_copy.trainable_variables))\n",
        "\n",
        "                # Saves regression line after 1st gradient step\n",
        "                if(i == 0):\n",
        "                    y_1grad_baseline = baseline_copy.f(x_truth)\n",
        "            \n",
        "            # Saves regression line after \"grad_steps\" gradient steps\n",
        "            y_ngrad_baseline = baseline_copy.f(x_truth)\n",
        "\n",
        "\n",
        "        # Plots\n",
        "\n",
        "        # MAML\n",
        "        # -----\n",
        "\n",
        "        fig = plt.figure()\n",
        "        plt.plot(x_truth, y_truth, 'r', x_truth, y_preupdate_maml, 'g:', \n",
        "                 x_truth, y_1grad_maml, 'b--', x_truth, y_ngrad_maml, 'm--', \n",
        "                 x_test, y_test, 'c^')\n",
        "        plt.legend([\"True Sine-Wave\", \"Pre-update\", \"1 Gradient Step\", \n",
        "                   f\"{grad_steps} Gradient Steps\", f\"{k} Points Used for Gradient\"], \n",
        "                   loc=\"upper center\", bbox_to_anchor= (0.5,-0.07), ncol=3)\n",
        "        plt.title(f\"MAML, K ={k}\")\n",
        "        fig.show()\n",
        "        fig.savefig(f\"MAML{k}.pdf\", bbox_inches='tight', pad_inches=0.25)\n",
        "        files.download(f\"MAML{k}.pdf\")\n",
        "\n",
        "        # Baseline\n",
        "        # ---------\n",
        "        \n",
        "        if(baseline):\n",
        "            fig = plt.figure()\n",
        "            plt.plot(x_truth, y_truth, 'r', x_truth, y_preupdate_baseline, 'g:', \n",
        "                    x_truth, y_1grad_baseline, 'b--', x_truth, y_ngrad_baseline, 'm--', \n",
        "                    x_test, y_test, 'c^')\n",
        "            plt.legend([\"True Sine-Wave\", \"Pre-update\", \"1 Gradient Step\", \n",
        "                    f\"{grad_steps} Gradient Steps\", f\"{k} Points Used for Gradient\"], \n",
        "                    loc=\"upper center\", bbox_to_anchor= (0.5,-0.07), ncol=3)\n",
        "            plt.title(f\"Baseline, K ={k}\")\n",
        "            fig.show()\n",
        "            fig.savefig(f\"Baseline{k}.pdf\", bbox_inches='tight', pad_inches=0.25)\n",
        "            files.download(f\"Baseline{k}.pdf\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NBAsbJrfl4XR"
      },
      "source": [
        "## Basline Class\n",
        "*Subclass of the MAML Class*\n",
        "\n",
        "---\n",
        "\n",
        "This class will create a baseline to test the MAML model against"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X0f0Nqsvl5Fw"
      },
      "source": [
        "class Baseline(MAML):\n",
        "    def __init__(self, k=100):\n",
        "        super().__init__(k)\n",
        "    \n",
        "    def train(self, epochs=10000, tasks=25):\n",
        "        num_epochs = epochs  \n",
        "        num_tasks = tasks   # Number of tasks per epoch\n",
        "        total_loss = []     # Log of losses for graphing\n",
        "\n",
        "\n",
        "        for epoch in range(num_epochs):\n",
        "            print(epoch)    # Track progress\n",
        "\n",
        "            # Sample batch of tasks T_i ~ p(T)\n",
        "            sine_tasks = [SineWave() for _ in range(num_tasks)]\n",
        "\n",
        "            # For each T_i do:\n",
        "            for sinusoid in sine_tasks:\n",
        "                # Sample K datapoints D from T_i\n",
        "                x_train, y_train = sinusoid.sample(self.k)\n",
        "\n",
        "                # Evaluate gradient w.r.t theta of Loss(f_theta)\n",
        "                with tf.GradientTape() as innerTape:\n",
        "                    y_hat = self.f(x_train)\n",
        "                    loss = tf.reduce_mean(self.loss_fnc(y_train, y_hat))\n",
        "                total_loss.append(loss)\n",
        "                gradient = innerTape.gradient(loss, self.trainable_variables)\n",
        "\n",
        "                # Update weights\n",
        "                self.optimizer.apply_gradients(zip(gradient, self.trainable_variables))\n",
        "\n",
        "\n",
        "        # Plots loss over epochs\n",
        "        fig = plt.figure()\n",
        "        plt.plot(range(epochs*tasks), total_loss)\n",
        "        plt.title(\"Loss per Epoch\")\n",
        "        fig.show()\n",
        "        fig.savefig(\"Baseline_Loss.pdf\", bbox_inches='tight', pad_inches=0.25)\n",
        "        files.download(\"Baseline_Loss.pdf\")"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}