example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "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.7.5"
    },
    "pycharm": {
      "stem_cell": {
        "cell_type": "raw",
        "metadata": {
          "collapsed": false
        },
        "source": []
      }
    },
    "colab": {
      "name": "example.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/dbaranchuk/memory-efficient-maml/blob/master/example.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6F74pLM66dNt",
        "colab_type": "text"
      },
      "source": [
        "## Gradient Checkpointing Model-Agnostic Meta-Learning\n",
        "\n",
        "We demonstrate how to use memory efficient MAML on CIFAR10.\n",
        "This notebook performs one forward and backward for MAML with a large number of iterations\n",
        "\n",
        "* Data: Random tensors (batch_size, 3, 224, 224)  \n",
        "* Model: ResNet18\n",
        "* Optimizer: SGD with 0.01 learning rate\n",
        "* Batch size: 16\n",
        "* MAML steps: 100 (works with >500 on 11GB GPU)\n",
        "* GPU: whatever colab has to spare, probably K80"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "pycharm": {
          "is_executing": true
        },
        "id": "alrX0irZ6dNv",
        "colab_type": "code",
        "outputId": "ef3005b7-9f87-4570-c662-09c829e4a06a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 289
        }
      },
      "source": [
        "%env CUDA_VISIBLE_DEVICES=0\n",
        "# colab dependencies\n",
        "!pip install torch==1.3.1 torchvision==0.4.2 torch_maml\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "import torch, torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision.models as models\n",
        "\n",
        "import torch_maml\n",
        "\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "\n",
        "# For reproducibility\n",
        "import random\n",
        "random.seed(42)\n",
        "np.random.seed(42)\n",
        "torch.manual_seed(42)\n",
        "\n",
        "torch.backends.cudnn.deterministic = True\n",
        "torch.backends.cudnn.benchmarks = False"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "env: CUDA_VISIBLE_DEVICES=0\n",
            "Requirement already satisfied: torch==1.3.1 in /usr/local/lib/python3.6/dist-packages (1.3.1)\n",
            "Requirement already satisfied: torchvision==0.4.2 in /usr/local/lib/python3.6/dist-packages (0.4.2)\n",
            "Collecting torch_maml\n",
            "  Downloading https://files.pythonhosted.org/packages/be/4c/a37a23fe88d41a47589e7653b398762a71d98d7dff8b2111759cc1a173e0/torch_maml-1.0.tar.gz\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch==1.3.1) (1.17.4)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from torchvision==0.4.2) (1.12.0)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision==0.4.2) (4.3.0)\n",
            "Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from pillow>=4.1.1->torchvision==0.4.2) (0.46)\n",
            "Building wheels for collected packages: torch-maml\n",
            "  Building wheel for torch-maml (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for torch-maml: filename=torch_maml-1.0-cp36-none-any.whl size=9396 sha256=4e6f09e990198a915667d462af6b5a50c0088c153a75144003320df25f071cba\n",
            "  Stored in directory: /root/.cache/pip/wheels/79/67/b2/923f59310ddb7a8de189573c3322a1af7754659ee472081bcc\n",
            "Successfully built torch-maml\n",
            "Installing collected packages: torch-maml\n",
            "Successfully installed torch-maml-1.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mx7mOCc76dNz",
        "colab_type": "text"
      },
      "source": [
        "#### Define compute_loss function and create model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EOvFGH9W6dN0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Interface:\n",
        "# def compute_loss(model, data, **kwargs):\n",
        "#      <YOUR CODE HERE>  # ideally this should be stateless (does not change global variables)\n",
        "#      return loss\n",
        "\n",
        "# Our example\n",
        "def compute_loss(model, data, device='cuda'):\n",
        "    inputs, targets = data\n",
        "    preds = model(inputs.to(device=device))\n",
        "    loss = F.cross_entropy(preds, targets.to(device=device))\n",
        "    return loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WwFkRoBP6dN2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Model is a torch.nn.Module \n",
        "model = models.resnet18(num_classes=10).to(device)\n",
        "# Optimizer is a custom MAML optimizer, e.g. SGD\n",
        "optimizer = torch_maml.IngraphGradientDescent(learning_rate=0.01)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cti7QScs6dN4",
        "colab_type": "text"
      },
      "source": [
        "#### Create NaiveMAML and GradientCheckpointMAML for comparison"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "35otHsPc6dN4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "efficient_maml = torch_maml.GradientCheckpointMAML(\n",
        "    model, compute_loss, optimizer=optimizer, checkpoint_steps=5)\n",
        "\n",
        "naive_maml = torch_maml.NaiveMAML(model, compute_loss, optimizer=optimizer)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nQxlYGRR6dN6",
        "colab_type": "text"
      },
      "source": [
        "#### Sanity check: small number of steps\n",
        "\n",
        "Both naive and memory-efficient maml should produce the same output."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I9KfEdYL6dN6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# First, we set such max steps that fits memory for naive MAML to check the implementation\n",
        "maml_steps = 10\n",
        "\n",
        "# Clip meta-learning gradients by global norm to avoid explosion\n",
        "max_grad_grad_norm = 1e2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZpSyUhuk6dN9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Generate batch for demonstration. Note that we support using different batches for each MAML step (a-la SGD)\n",
        "x_batch, y_batch = torch.randn((16, 3, 224, 224)), torch.randint(0, 10, (16, ))\n",
        "inputs = [(x_batch, y_batch)] * maml_steps  # use the same batch for each step"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T5hbRPIM6dN_",
        "colab_type": "code",
        "outputId": "3289b5b5-728f-47b3-8d83-b3ac06b1552d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "updated_model, loss_history, _ = naive_maml(inputs, loss_kwargs={'device':device},\n",
        "                                            max_grad_grad_norm=max_grad_grad_norm)\n",
        "final_loss = compute_loss(updated_model, (x_batch, y_batch), device=device)\n",
        "final_loss.backward()\n",
        "grads_naive = [params.grad for params in model.parameters()]\n",
        "print(\"Loss naive: %.4f\" % final_loss.item())"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loss naive: 0.5553\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j2G812Bs6dOA",
        "colab_type": "code",
        "outputId": "80951ba4-03a1-4007-9300-02e23a74a0e5",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "updated_model, loss_history, _ = efficient_maml(inputs, loss_kwargs={'device':device},\n",
        "                                                max_grad_grad_norm=max_grad_grad_norm)\n",
        "final_loss = compute_loss(updated_model, (x_batch, y_batch), device=device)\n",
        "final_loss.backward()\n",
        "grads_efficient = [params.grad for params in model.parameters()]\n",
        "print(\"Loss memory-efficient: %.4f\" % final_loss.item())"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loss memory-efficient: 0.5553\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lTnR_KFM6dOC",
        "colab_type": "code",
        "outputId": "5de0dbbc-c3bf-4a67-8c52-bc989005df78",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "for grad1, grad2 in zip(grads_naive, grads_efficient):\n",
        "    assert torch.allclose(grad1, grad2)\n",
        "\n",
        "print(\"All grads match!\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "All grads match!\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8XfRGRhq6dOD",
        "colab_type": "code",
        "outputId": "d618ef2c-03b0-48d9-ebfd-5f8adffabdd1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "# alternative: use rmsprop optimizer\n",
        "rmsprop_maml = torch_maml.GradientCheckpointMAML(\n",
        "    model, compute_loss, optimizer=torch_maml.IngraphRMSProp(learning_rate=1e-3, beta=0.9, epsilon=1e-5), \n",
        "    checkpoint_steps=5)\n",
        "\n",
        "updated_model, loss_history, _ = rmsprop_maml(inputs, loss_kwargs={'device':device},\n",
        "                                                max_grad_grad_norm=max_grad_grad_norm)\n",
        "final_loss = compute_loss(updated_model, (x_batch, y_batch), device=device)\n",
        "final_loss.backward()\n",
        "grads_efficient = [params.grad for params in model.parameters()]\n",
        "print(\"Loss RMSProp: %.4f\" % final_loss.item())"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loss RMSProp: 0.0224\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t13US_Vn6dOF",
        "colab_type": "text"
      },
      "source": [
        "### The real meta-learning: 100 steps and beyond"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gUk7JFIm6dOG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "maml_steps = 100  # feel free to tweak (works with >500)\n",
        "inputs = [(x_batch, y_batch)] * maml_steps\n",
        "torch.cuda.empty_cache()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qx4yJPAJ6dOH",
        "colab_type": "code",
        "outputId": "756a11ba-424f-4709-e384-c8177c99663e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 284
        }
      },
      "source": [
        "updated_model, loss_history, _ = efficient_maml(inputs, loss_kwargs={'device':device},\n",
        "                                        max_grad_grad_norm=max_grad_grad_norm)\n",
        "final_loss = compute_loss(updated_model, (x_batch, y_batch), device=device)\n",
        "final_loss.backward()\n",
        "grads_efficient = [params.grad for params in model.parameters()]\n",
        "\n",
        "plt.plot(loss_history)\n",
        "print(\"Loss memory-efficient: %.4f\" % final_loss.item())"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loss memory-efficient: 0.0427\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAAD6CAYAAACxrrxPAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjIsIGh0\ndHA6Ly9tYXRwbG90bGliLm9yZy8li6FKAAAcZ0lEQVR4nO3deZRc5Xnn8e9TS1f13i11q7VLCHUQ\nAoOAHhkGY7BhEnA81syEGaOZxI7HHsWOfbzOeJzMxD5x/pg48XESBg6EYMcmdrBj8Bjs4IVg2cbY\nBlrsSAgJIZBES2ptvamXWp75o24varXUJam6b9et3+ecOnWXt7ueyxW/uv3e995r7o6IiJS/WNgF\niIhIaSjQRUQiQoEuIhIRCnQRkYhQoIuIRIQCXUQkIqYNdDNbZmabzWyrmb1oZh+bos11ZtZjZs8E\nr8/OTLkiInIqiSLaZIFPuftTZlYPbDGzh91966R2j7r7O4v94JaWFl+5cuUZlCoiIlu2bDnk7q1T\nrZs20N29C+gKpvvMbBuwBJgc6Gdk5cqVdHZ2nsuvEBGpOGb22qnWnVEfupmtBC4DHp9i9VVm9qyZ\n/cDMLjqjCkVE5JwV0+UCgJnVAfcDH3f33kmrnwJWuHu/mb0D+C7QPsXv2ARsAli+fPlZFy0iIicr\n6gjdzJIUwvwb7v6dyevdvdfd+4Pph4CkmbVM0e4ud+9w947W1im7gERE5CwVM8rFgC8D29z9S6do\nszBoh5mtD37v4VIWKiIip1dMl8vVwO8Bz5vZM8GyPwaWA7j7ncDNwIfMLAsMAre4buMoIjKrihnl\n8gvApmlzG3BbqYoSEZEzpytFRUQiouwCffv+Pr74o+0cGRgJuxQRkTml7AJ9V3c/t23eyf6eobBL\nERGZU8ou0OvTSQD6hjIhVyIiMreUYaAXzuP2D2dDrkREZG4pu0CvCwK9b0iBLiIyUdkFev1YoKvL\nRURkovIL9FTQh64uFxGRE5RdoKeTMRIxU5eLiMgkZRfoZkZ9OqEuFxGRScou0KFwYrRfR+giIico\ny0CvTyXV5SIiMkl5Bno6oZOiIiKTlG+g6whdROQEZRroSZ0UFRGZpEwDPaFL/0VEJinLQK9LFbpc\n9FAkEZFxZRno9ekkubwzmMmFXYqIyJxRloE+eoMujUUXERlXloHeEAR6rwJdRGRMWQa67rgoInKy\nsgz0uuCOixrpIiIyriwDvV4PuRAROUlZB7pOioqIjCvPQA+6XHrVhy4iMqYsA13PFRUROVlZBno8\nZtRWxXVSVERkgrIMdCgcpWvYoojIuLIN9MIdF3WELiIyqowDXXdcFBGZqGwDvS6V0KX/IiITlG2g\nN6ST9KsPXURkTNkG+ug90UVEpGDaQDezZWa22cy2mtmLZvaxKdqYmd1qZjvN7Dkzu3xmyh2n54qK\niJyomCP0LPApd18LXAl82MzWTmpzE9AevDYBd5S0yinUp5MMZnJkc/mZ/igRkbIwbaC7e5e7PxVM\n9wHbgCWTmm0A7vGCXwNNZrao5NVOMPaQC410EREBzrAP3cxWApcBj09atQTYM2F+LyeHfknpjosi\nIicqOtDNrA64H/i4u/eezYeZ2SYz6zSzzu7u7rP5FWMaFOgiIicoKtDNLEkhzL/h7t+Zosk+YNmE\n+aXBshO4+13u3uHuHa2trWdT75jRh1zo8n8RkYJiRrkY8GVgm7t/6RTNHgTeE4x2uRLocfeuEtZ5\nEnW5iIicKFFEm6uB3wOeN7NngmV/DCwHcPc7gYeAdwA7gePA+0pf6onqdVJUROQE0wa6u/8CsGna\nOPDhUhVVjDo9KFpE5ARle6VoQzroQ9cRuogIUMaBnkrESMZNfegiIoGyDXQzC+7noi4XEREo40CH\nwuX//TpCFxEByjzQdcdFEZFxZR3ouuOiiMi4Mg/0pEa5iIgEyjzQdVJURGRUBAJdR+giIhCBQO8f\nzlK4UFVEpLKVdaDXpZLk8s5gJhd2KSIioSvrQB+7QZe6XUREohHovQp0EZFoBLpGuoiIlHmgN9dU\nAXBkYCTkSkREwlfWgb6wMQ3A/t6hkCsREQlfWQd6a12KmMGBHgW6iEhZB3oiHqO1PkWXAl1EpLwD\nHWBhQ1pdLiIiRCDQ2xrSHFCgi4iUf6AvakyzX10uIiLlH+htjWl6h7IcH9HFRSJS2co+0Bc2BEMX\ndZQuIhUuOoGufnQRqXDlH+iNOkIXEYEoBbqO0EWkwpV9oNdUJahPJ3S1qIhUvLIPdAiGLuoIXUQq\nXCQCva1BY9FFRCIR6Lr8X0QkIoG+qDFNd98w2Vw+7FJEREITiUBva0yTd+juHw67FBGR0EQi0HW1\nqIhIEYFuZl8xs4Nm9sIp1l9nZj1m9kzw+mzpyzy90bHouuuiiFSyRBFtvgrcBtxzmjaPuvs7S1LR\nWdARuohIEUfo7v5z4Mgs1HLW5tVWURWP0aUjdBGpYKXqQ7/KzJ41sx+Y2UWnamRmm8ys08w6u7u7\nS/TRYGa0NaZ0taiIVLRSBPpTwAp3vxT4v8B3T9XQ3e9y9w5372htbS3BR4/TWHQRqXTnHOju3uvu\n/cH0Q0DSzFrOubIzpKtFRaTSnXOgm9lCM7Ngen3wOw+f6+89U6NH6O4+2x8tIjInTDvKxczuBa4D\nWsxsL/A5IAng7ncCNwMfMrMsMAjc4iGk6sLGNEOZPL2DWRprkrP98SIioZs20N194zTrb6MwrDFU\nE++LrkAXkUoUiStFYXws+hs9gyFXIiISjsgE+or5tQC8dmgg5EpERMIRmUBvqauisTrJjoP9YZci\nIhKKyAS6mdG+oE6BLiIVKzKBDtDeVsdOBbqIVKhIBfrqBfUcGRjhsO6LLiIVKFKB3r6gDkBH6SJS\nkSIV6KuDQFc/uohUokgF+qLGNLVVcR2hi0hFilSgmxmr2+rZcbAv7FJERGZdpAIdCv3oOw7oCF1E\nKk8kA/1g3zA9g5mwSxERmVWRC/TVGukiIhUqcoHevqAegJ3qRxeRChO5QF/SXE06GVM/uohUnMgF\nejxmnN+qe7qISOWJXKBDoR9dfegiUmkiGejtC+rYd2yQgeFs2KWIiMyaSAb66uDE6CvdOkoXkcoR\nyUBvbysMXdy+XyNdRKRyRDLQV86vpT6V4Ok9x8IuRURk1kQy0OMx47IVzWzZfTTsUkREZk0kAx2g\nY0UzLx/s0y0ARKRiRDrQ3eGp13WULiKVIbKBvm55E/GY8dRrCnQRqQyRDfSaqgRrFzXQqX50EakQ\nkQ10gCtWNPPMnmNkcvmwSxERmXGRD/TBTI5tXb1hlyIiMuMiHegdK5sB1O0iIhUh0oG+qLGaJU3V\nbNGJURGpAJEOdCh0u3S+dgR3D7sUEZEZFflA71jZzIHeYfYdGwy7FBGRGTVtoJvZV8zsoJm9cIr1\nZma3mtlOM3vOzC4vfZln7/LlhX50dbuISNQVc4T+VeDG06y/CWgPXpuAO869rNK5cFEDjdVJfv7y\nobBLERGZUdMGurv/HDhymiYbgHu84NdAk5ktKlWB5yoeM952QSubtx8kl1c/uohEVyn60JcAeybM\n7w2WzRk3rG3jyMAIT+u+LiISYbN6UtTMNplZp5l1dnd3z9rnvvU3WknGjYe3HZi1zxQRmW2lCPR9\nwLIJ80uDZSdx97vcvcPdO1pbW0vw0cVpSCd583nz+ZetCnQRia5SBPqDwHuC0S5XAj3u3lWC31tS\nN1y4gFe6B3j10EDYpYiIzIhihi3eC/wKuMDM9prZ+83sg2b2waDJQ8AuYCfwd8Afzli15+D6C9sA\neETdLiISUYnpGrj7xmnWO/DhklU0Q5bNq2HNwnoe3nqAD1yzKuxyRERKLvJXik50w4VtdL52lGPH\nR8IuRUSk5Coq0K+/cAG5vLN5+8GwSxERKbmKCvRLlzbR1pDi+8/OuXO2IiLnrKICPRYz/v1lS/np\ny90c7BsKuxwRkZKqqEAHuPmKpeTyzgNPvxF2KSIiJVVxgb56QR2XLW/ivi17dY90EYmUigt0KByl\nbz/Qx/P7esIuRUSkZCoy0N95yWJSiRj3bdkbdikiIiVTkYHeWJ3kty5ayAPPvMFwNhd2OSIiJVGR\ngQ6FbpeewQyPbNOYdBGJhooN9KtXt7CoMc3Xf/1a2KWIiJRExQZ6PGa891+v5JevHObFN3RyVETK\nX8UGOsDG9cuprYpz96Ovhl2KiMg5q+hAb6xO8u5/tZzvPfsGXT2DYZcjInJOKjrQAd539Ury7nz1\nsd1hlyIick4qPtCXzavhpjct4h8ff52+oUzY5YiInLWKD3SATdesom84y7ee3BN2KSIiZ02BDly6\nrIkrV83jzp/tYmA4G3Y5IiJnRYEe+PSNazjUP6wRLyJSthTogcuXN3PTxQv525+/QnffcNjliIic\nMQX6BP/jty5gJJvn1kd2hF2KiMgZU6BPsKq1jo3rl/OPT7zOru7+sMsRETkjCvRJPnp9O+lEjD//\nwUthlyIickYU6JO01qf4w7et5sdbD7B5u+7EKCLlQ4E+hf92zSpWtdbyuQdeZCij+6WLSHlQoE+h\nKhHjzzZczOtHjnPHT18JuxwRkaIo0E/h6tUtvOvSxdzxs1fYfWgg7HJERKalQD+N//3bF1IVj/En\nD7xAPu9hlyMicloK9NNY0JDmf960hkd3HOLvf7k77HJERE5LgT6N333zcm64sI0v/OAlPdlIROY0\nBfo0zIy/uPkSmmqSfPTepzk+opt3icjcpEAvwrzaKv7q3evYdWiAz39va9jliIhMSYFepKtXt/DB\na8/nm0/u4d4nXg+7HBGRkxQV6GZ2o5ltN7OdZvaZKdb/vpl1m9kzwesDpS81fP/9Ny/gmvYWPvvA\nC3TuPhJ2OSIiJ5g20M0sDtwO3ASsBTaa2dopmn7L3dcFr7tLXOecEI8Zt228nCVN1Xzw60/pwdIi\nMqcUc4S+Htjp7rvcfQT4JrBhZsuauxprkvzdezoYHMmy6Z4tesKRiMwZxQT6EmDiwzb3Bssm+x0z\ne87M7jOzZSWpbo5qb6vn1o2XsbWrlw9+fQvDWd3vRUTCV6qTot8DVrr7JcDDwNemamRmm8ys08w6\nu7u7S/TR4bj+wjb+z394E4/uOMQn/+lZcrqSVERCVkyg7wMmHnEvDZaNcffD7j763La7gSum+kXu\nfpe7d7h7R2tr69nUO6f8p45l/NFNa/jn57r43IMv4K5QF5HwJIpo8yTQbmbnUQjyW4D/PLGBmS1y\n965g9l3AtpJWOYf9wbXnc+T4CH/7s124w59tuJhYzMIuS0Qq0LSB7u5ZM/sI8CMgDnzF3V80s88D\nne7+IPBRM3sXkAWOAL8/gzXPOZ+5cQ2GcefPXmFwJMdf3HwJibiG+IvI7LKwugk6Ojq8s7MzlM+e\nCe7O7Zt38sUfv8yNFy3kr29ZRzoZD7ssEYkYM9vi7h1TrdNhZImYGR95ezt/8s61/PDF/dxy1685\n2DsUdlkiUkEU6CX2/recx52/ewXb9/ex4fbHeGGf7tAoIrNDgT4Dbrx4Ifd96CoMuPnOX3L/lr1h\nlyQiFUCBPkMuWtzIAx95C+uWNfGpbz/Lp+97lsERXYAkIjNHgT6DWutTfP39b+ajb1/Nt7fsZcPt\nv9BDMkRkxijQZ1giHuOTv3kB9/zX9RwZyLDhtsf40sMvM5LNh12aiESMAn2WXNPeysOfeCv/9tLF\n3PrIDt512y/Y8trRsMsSkQhRoM+i5uDJR19+bwfHjmf4nTt+ySf/6RkNbxSRkijm0n8psesvbOPK\nVfO5ffNO7n70VX70wn4+dN35vO/q86hNaZeIyNnREXpIalMJPn3jGn78ibdy9eoWvvjjl7n2Lzfz\n1cde1e14ReSs6NL/OeKp14/ylz/czq92HWZhQ5o/uHYVG9cv1+0DROQEp7v0X4E+h7g7j+08zK0/\n2cETrx6hpS7F+65eycb1y5lXWxV2eSIyByjQy9Djuw5z2+adPLrjEFWJGBsuXcx7rlrJm5Y2hl2a\niITodIGuM3Bz1JtXzefNq+az40AfX/vVbu7fso9vb9nLRYsbuGX9cjasW0xDOhl2mSIyh+gIvUz0\nDGZ44Jl93PvEHrZ19VKViHH9mgVsWLeEt61pJZVQX7tIJVCXS4S4O8/t7eH/Pb2P7z/3Bof6R6hP\nJbhhbRvveNMirmlv0YlUkQhToEdUNpfnsVcO88/PvcGPXjxAz2CG2qo4b2lv4fo1bVy3ppUF9emw\nyxSRElKgV4BMLs9jOw/x8NYD/OSlg3T1FK4+vWhxA9f+RivXtLdy+Yomdc2IlDkFeoVxd7Z29fLT\n7d387OVunnrtKNm8k0rEuGJFM1etms/68+Zx6bImdc+IlBkFeoXrHcrw+K4j/OqVw/xq12G2dfUC\nUBWPccnSRi5f0cxly5pYt7yJRY3VIVcrIqejQJcTHDs+Qufuozyx+whP7j7Ci/t6GckVbufbWp/i\nkiWNvGlpIxctbmTt4gYWN6Yxs5CrFhHQOHSZpKmmihvWtnHD2jYAhrM5tnX18fTrR3l+Xw/P7+3h\nJ9sPMvpd31id5IKF9axZWM8FC+u5oK2e9gX1NNZoHLzIXKJAF1KJOOuWNbFuWdPYsoHhLC/t72Nr\nVy9b3+hl+/5e7t+yl4EJj9FrrU9xfmst57fWsaq1jlUttaxsqWVpczXJuO77JjLbFOgypdpUgitW\nNHPFiuaxZfm8s+/YIDsP9rPjYB8vH+hnV3c/33+ui57BzFi7eMxY2lzN8nk1rJhfw4p5hZBfNq+G\nZc01NFQn1IUjMgMU6FK0WMwKoTyvhretWTC23N05PDDC7kMDvHpogN2HB9h9+Dh7jhzne8+eGPYA\ndakES5qqWdyUZnFTNYubqlnUmGZhY5qFDWkWNVZTXaXRNyJnSoEu58zMaKlL0VKXomPlvJPW9wxm\n2HPkOHuPHmfv0UH2Hh1k37FB9h0d5Jk9xzh6PHPSz9SnEyxsSNPWkGZBfYrWhhQL6tO01FXRWpei\npb7weU3VSWIxHe2LgAJdZkFjdZLGJY1cvGTqO0UOjuR4o2eQAz1D7O8doqtniIO9QxzoHWZ/7xCv\nHhqgu294bCTORPGYMa+2ivm1Vcyvq2J+bYp5tVVjr+aaKpprkjTXVtFUk6S5pkpj7yWyFOgSuuqq\nOOe31nF+a90p27g7x45nONQ/THf/MN19wxwZGOFw/wiH+oc5PDDC4f5hnj16jCP9I/QNZ0/5u9LJ\nGE3VVYUvmppk4X3CqyGdoKE6SUM6SX0wXZ9OUJ9OUpdKENdfBDJHKdClLJgZzbVVNNdW0d5WP237\nkWyeo8dHCq+BDEePj3Ds+Oj7CD2DGY4dz3As6A56cTBDz2DmhFE8p1KXShRe6QT16cTYfO3Ye3xs\nuqYqQV0qTk1VYX1tKk5tVYKaqsKydDKmE8RSMgp0iaSqRIy2oA/+TGRyefqGsvQGAd83lKVvKEPv\nUGG6N5gfGM7SP5ylb6jwfqB3iP5gemAkRy5f3AV7ZlCTjFM9FvJxqkffk4nCdLKwLJ2MU52Mk07G\nxuYnLksn46QT49OpZIxUMF8V1xdHJVCgi0yQjMfG+t/PlrsznM0Xwj0I/sGRHAMjOQaGsxwfyXF8\nZHz56PzxkRyDIzkGM4X3owODY9Oj71OdRyiGGaQToyE/HvSpRLwwH4R+KmhTFR9dNj5flSj8bFVi\nfH50OpmIkQreq+IxkvET2yXjFrwXlulE9sxQoIuUmJmNHT231KVK+rtzeWcoMx7ww9kcgyN5hrI5\nhjI5hjL54D3HUDbPcCbHcHZ82XA2z3Cm0H4kmx9bN5zJ0zuYZXjC8uFsnpHR11l+kZxKPGYk40Yy\nCP/J04nY6JfD+HQyVmiTCNomYkZiYvu4kThhOmgzqd1om9Ea4sHvjU9oW3gvzMdjseDdxt/j48tG\nl8+Fv4AU6CJlJB6zoC9+dv/XzeedkVwh2EdDPhNMD2fHl2cmvA9n82RyHrTPkc372HwmN94mm8+T\nyXphWd7JjK7PO9mg3eBgbuxnsjknkw/ec052bDpPNu9Fd3eVWswY+6JIxIx43IjbeODHJnwhbFy/\nnA9cs6rkNRT1r8LMbgT+BogDd7v7n09anwLuAa4ADgPvdvfdpS1VRMISixnpWLwshnzm8042Xwj6\nTK4Q8NngCyIXfBnk8j725ZAN1ucm/FwuD7mJPx+0Gf3CyOTy5H10eWFZ3gtfMHn34PfmT1yed3LB\nz5T6L7dR0wa6mcWB24F/A+wFnjSzB91964Rm7weOuvtqM7sF+ALw7pkoWETkdGIxoypmVFF59xMq\nZovXAzvdfZe7jwDfBDZMarMB+FowfR9wvc2FDiURkQpSTKAvAfZMmN8bLJuyjbtngR5gfikKFBGR\n4szq3yRmtsnMOs2ss7u7ezY/WkQk8ooJ9H3AsgnzS4NlU7YxswTQSOHk6Anc/S5373D3jtbW1rOr\nWEREplRMoD8JtJvZeWZWBdwCPDipzYPAe4Ppm4GfeFjPthMRqVDTjnJx96yZfQT4EYVhi19x9xfN\n7PNAp7s/CHwZ+Acz2wkcoRD6IiIyi4oah+7uDwEPTVr22QnTQ8B/LG1pIiJyJipvoKaISERZWF3d\nZtYNvHaWP94CHCphOeWiEre7ErcZKnO7K3Gb4cy3e4W7TzmqJLRAPxdm1unuHWHXMdsqcbsrcZuh\nMre7ErcZSrvd6nIREYkIBbqISESUa6DfFXYBIanE7a7EbYbK3O5K3GYo4XaXZR+6iIicrFyP0EVE\nZJKyC3Qzu9HMtpvZTjP7TNj1zAQzW2Zmm81sq5m9aGYfC5bPM7OHzWxH8N4cdq0zwcziZva0mX0/\nmD/PzB4P9vm3gltQRIaZNZnZfWb2kpltM7OrKmFfm9kngn/fL5jZvWaWjuK+NrOvmNlBM3thwrIp\n968V3Bps/3NmdvmZfFZZBfqEh23cBKwFNprZ2nCrmhFZ4FPuvha4EvhwsJ2fAR5x93bgkWA+ij4G\nbJsw/wXgr9x9NXCUwgNVouRvgB+6+xrgUgrbHul9bWZLgI8CHe5+MYXbiow+HCdq+/qrwI2Tlp1q\n/94EtAevTcAdZ/JBZRXoFPewjbLn7l3u/lQw3Ufhf/AlnPggka8B/y6cCmeOmS0Ffhu4O5g34O0U\nHpwCEdtuM2sE3krhfki4+4i7H6MC9jWFW49UB3dorQG6iOC+dvefU7jH1USn2r8bgHu84NdAk5kt\nKvazyi3Qi3nYRqSY2UrgMuBxoM3du4JV+4G2kMqaSX8NfBoYfcz8fOBY8OAUiN4+Pw/oBv4+6Ga6\n28xqifi+dvd9wBeB1ykEeQ+whWjv64lOtX/PKePKLdAripnVAfcDH3f33onrgtsTR2qIkpm9Ezjo\n7lvCrmUWJYDLgTvc/TJggEndKxHd180UjkbPAxYDtZzcLVERSrl/yy3Qi3nYRiSYWZJCmH/D3b8T\nLD4w+udX8H4wrPpmyNXAu8xsN4XutLdT6F9uCv4sh+jt873AXnd/PJi/j0LAR31f3wC86u7d7p4B\nvkNh/0d5X090qv17ThlXboFezMM2yl7Qb/xlYJu7f2nCqokPEnkv8MBs1zaT3P2P3H2pu6+ksG9/\n4u7/BdhM4cEpELHtdvf9wB4zuyBYdD2wlYjvawpdLVeaWU3w7310uyO7ryc51f59EHhPMNrlSqBn\nQtfM9Ny9rF7AO4CXgVeA/xV2PTO0jW+h8CfYc8AzwesdFPqTHwF2AP8CzAu71hn8b3Ad8P1gehXw\nBLAT+DaQCru+Em/rOqAz2N/fBZorYV8Dfwq8BLwA/AOQiuK+Bu6lcJ4gQ+Evsvefav8CRmEk3yvA\n8xRGARX9WbpSVEQkIsqty0VERE5BgS4iEhEKdBGRiFCgi4hEhAJdRCQiFOgiIhGhQBcRiQgFuohI\nRPx/W1Az+EYhq14AAAAASUVORK5CYII=\n",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3OGePjtF6dOI",
        "colab_type": "code",
        "outputId": "0a17be53-8bc9-4f9c-b7f1-d95291060dd9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 375
        }
      },
      "source": [
        "# naive maml can't handle this...\n",
        "updated_model, loss_history, _ = naive_maml(inputs, loss_kwargs={'device':device},\n",
        "                                            max_grad_grad_norm=max_grad_grad_norm)\n",
        "final_loss = compute_loss(updated_model, (x_batch, y_batch), device=device)\n",
        "final_loss.backward()\n",
        "grads_naive = [params.grad for params in model.parameters()]\n",
        "print(\"Loss naive: %.4f\" % final_loss.item())"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "error",
          "ename": "RuntimeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-13-9aa961209633>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m updated_model, loss_history, _ = naive_maml(inputs, loss_kwargs={'device':device},\n\u001b[0;32m----> 2\u001b[0;31m                                             max_grad_grad_norm=max_grad_grad_norm)\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0mfinal_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompute_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mupdated_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mx_batch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_batch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mfinal_loss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mgrads_naive\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrad\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mparams\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparameters\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    539\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    540\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 541\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    542\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    543\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch_maml/maml.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, inputs, opt_kwargs, loss_kwargs, optimizer_state, **kwargs)\u001b[0m\n\u001b[1;32m     51\u001b[0m         \u001b[0mloss_history\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     52\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0minput\u001b[0m \u001b[0;32min\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 53\u001b[0;31m             \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mupdated_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mloss_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     54\u001b[0m             \u001b[0mloss_history\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     55\u001b[0m             optimizer_state, updated_model = self.optimizer.step(\n",
            "\u001b[0;32m<ipython-input-2-23aa1756c359>\u001b[0m in \u001b[0;36mcompute_loss\u001b[0;34m(model, data, device)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mcompute_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'cuda'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m     \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m     \u001b[0mpreds\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m     \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcross_entropy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpreds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    539\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    540\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 541\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    542\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    543\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torchvision/models/resnet.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    194\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    195\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 196\u001b[0;31m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    197\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbn1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    198\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrelu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    539\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    540\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 541\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    542\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    543\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    343\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    344\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 345\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv2d_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    346\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    347\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mConv3d\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_ConvNd\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py\u001b[0m in \u001b[0;36mconv2d_forward\u001b[0;34m(self, input, weight)\u001b[0m\n\u001b[1;32m    340\u001b[0m                             _pair(0), self.dilation, self.groups)\n\u001b[1;32m    341\u001b[0m         return F.conv2d(input, weight, self.bias, self.stride,\n\u001b[0;32m--> 342\u001b[0;31m                         self.padding, self.dilation, self.groups)\n\u001b[0m\u001b[1;32m    343\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    344\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 50.00 MiB (GPU 0; 15.90 GiB total capacity; 14.81 GiB already allocated; 33.88 MiB free; 362.32 MiB cached)"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3Hv17vTo6dOJ",
        "colab_type": "text"
      },
      "source": [
        "Now before you apply it, take a look at \"Tips and Tricks\" section in the README."
      ]
    }
  ]
}