EWC_MultiTaskLearning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "4VuHu3XXgCTo"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchvision.datasets as datasets\n",
        "import torchvision.transforms as transforms\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/ContinualAI/colab.git continualai/colab\n",
        "\n",
        "from continualai.colab.scripts import mnist\n",
        "mnist.init()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V04-pnYNgIFC",
        "outputId": "124d9ebe-1b4c-4d6a-a829-bacb8aa9660f"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'continualai/colab'...\n",
            "remote: Enumerating objects: 378, done.\u001b[K\n",
            "remote: Counting objects: 100% (120/120), done.\u001b[K\n",
            "remote: Compressing objects: 100% (58/58), done.\u001b[K\n",
            "remote: Total 378 (delta 78), reused 64 (delta 62), pack-reused 258\u001b[K\n",
            "Receiving objects: 100% (378/378), 26.97 MiB | 28.09 MiB/s, done.\n",
            "Resolving deltas: 100% (198/198), done.\n",
            "Downloading train-images-idx3-ubyte.gz...\n",
            "Downloading t10k-images-idx3-ubyte.gz...\n",
            "Downloading train-labels-idx1-ubyte.gz...\n",
            "Downloading t10k-labels-idx1-ubyte.gz...\n",
            "Download complete.\n",
            "Save complete.\n",
            "Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz\n",
            "Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to data/mnist/MNIST/raw/train-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 9912422/9912422 [00:00<00:00, 35392198.65it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting data/mnist/MNIST/raw/train-images-idx3-ubyte.gz to data/mnist/MNIST/raw\n",
            "\n",
            "Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz\n",
            "Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz to data/mnist/MNIST/raw/train-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28881/28881 [00:00<00:00, 1208637.50it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting data/mnist/MNIST/raw/train-labels-idx1-ubyte.gz to data/mnist/MNIST/raw\n",
            "\n",
            "Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz\n",
            "Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz to data/mnist/MNIST/raw/t10k-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 1648877/1648877 [00:00<00:00, 9477744.76it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting data/mnist/MNIST/raw/t10k-images-idx3-ubyte.gz to data/mnist/MNIST/raw\n",
            "\n",
            "Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz\n",
            "Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz to data/mnist/MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 4542/4542 [00:00<00:00, 1175160.62it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting data/mnist/MNIST/raw/t10k-labels-idx1-ubyte.gz to data/mnist/MNIST/raw\n",
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x_train, t_train, x_test, t_test = mnist.load()\n",
        "\n",
        "print(\"x_train dim and type: \", x_train.shape, x_train.dtype)\n",
        "print(\"t_train dim and type: \", t_train.shape, t_train.dtype)\n",
        "print(\"x_test dim and type: \", x_test.shape, x_test.dtype)\n",
        "print(\"t_test dim and type: \", t_test.shape, t_test.dtype)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OtZd26Q9gOmt",
        "outputId": "d634368d-b619-4c24-dbab-37ff52a73a62"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "x_train dim and type:  (60000, 1, 28, 28) float32\n",
            "t_train dim and type:  (60000,) uint8\n",
            "x_test dim and type:  (10000, 1, 28, 28) float32\n",
            "t_test dim and type:  (10000,) uint8\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "f, axarr = plt.subplots(2, 2)\n",
        "axarr[0, 0].imshow(x_train[1, 0], cmap=\"gray\")\n",
        "axarr[0, 1].imshow(x_train[2, 0], cmap=\"gray\")\n",
        "axarr[1, 0].imshow(x_train[3, 0], cmap=\"gray\")\n",
        "axarr[1, 1].imshow(x_train[4, 0], cmap=\"gray\")\n",
        "np.vectorize(lambda ax: ax.axis('off'))(axarr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 542
        },
        "id": "qBUrh26PgVDZ",
        "outputId": "125ae3aa-3754-4575-c451-5110ff872346"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(array([[-0.5, -0.5],\n",
              "        [-0.5, -0.5]]),\n",
              " array([[27.5, 27.5],\n",
              "        [27.5, 27.5]]),\n",
              " array([[27.5, 27.5],\n",
              "        [27.5, 27.5]]),\n",
              " array([[-0.5, -0.5],\n",
              "        [-0.5, -0.5]]))"
            ]
          },
          "metadata": {},
          "execution_count": 4
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 4 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAcoAAAGFCAYAAAB9krNlAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAARdElEQVR4nO3daYiVdfsHcI9JZoW2WRbRbkGFTZltREaJxVNUKLSgmb2oSIqICkmmKNpXsMiKxBYSbDFTi7Aos8UUbYMW2ynMmGwxtU1izv/FE/x5eq77es7tnDMzZ87n8/LL4Z4rOd1ff3D5m0q1Wq32AwBC/Xt6AADozRQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQG1PrBSqXSyDmgS9yb0Xt5d9Cb1fLucKIEgISiBICEogSAhKIEgISiBICEogSAhKIEgISiBICEogSAhKIEgISiBICEogSAhKIEgISiBICEogSAhKIEgISiBICEogSAhKIEgISiBICEogSAhKIEgISiBIDEgJ4egP80cuTIML/kkkvCfNKkSWH+2GOPhfm9994b5u+8804N0wG0HidKAEgoSgBIKEoASChKAEgoSgBIVKrVarWmD1YqjZ6lpbS1tYX5K6+8EuaDBw+uy8/95ZdfwnzHHXesy/N7So1fY3qAd0ffduKJJ4b57Nmzw3z06NFh/sknn9RtpjJqeXc4UQJAQlECQEJRAkBCUQJAQlECQMJdrw12xBFHhPncuXPDfMiQIWFetJm1YcOGMN+0aVOYF223HnXUUWFedAds0fOhux133HFhXvRdnzdvXiPHaTmjRo0K8xUrVnTzJI3jRAkACUUJAAlFCQAJRQkACUUJAAlbryVtvfXWYX7YYYeF+eOPPx7mu+66a13m+eyzz8L89ttvD/M5c+aE+Ztvvhnm7e3tYX7LLbfUMB003vHHHx/mw4cPD3Nbr5unf//4XLX33nuH+Z577hnmzXj3rxMlACQUJQAkFCUAJBQlACQUJQAkbL2W9OCDD4b5Oeec082T/FvRtu22224b5kuWLAnzos3BESNGbNZc0F0mTZoU5m+99VY3T9K3FW3qX3DBBWFetPG/atWqus3UXZwoASChKAEgoSgBIKEoASChKAEgYeu1wMiRI8P8lFNOCfOy9xcWbZ8uXLgwzO+8884wX7NmTZi/++67Yf7zzz+H+QknnBDmzXgvI62l6A5S6mvmzJmlPl90D3Uz8g0DgISiBICEogSAhKIEgISiBIBEy2+9trW1hflLL70U5oMHDw7zarUa5i+88EKYF90NO3r06DBvb28P86JNtLVr14b5+++/H+adnZ1hXrTlW3TH7DvvvBPm0FVF9w7vsssu3TxJaxoyZEipzxe9Q5uREyUAJBQlACQUJQAkFCUAJBQlACRaZut1//33D/OrrroqzIs2vH744Ycw/+6778L80UcfDfONGzeG+fPPP18qb7RBgwaF+RVXXBHmEyZMaOQ4tLB//etfYV70HWXzFG0R77333qWe8+2339ZjnF7BiRIAEooSABKKEgASihIAEooSABJ9but14MCBYX7nnXeGedEm3YYNG8J80qRJYb5y5cow76sbeXvssUdPj0CLOeCAA0p9/sMPP2zQJH1b0buyaBv2008/DfOid2gzcqIEgISiBICEogSAhKIEgISiBIBEn9t6PfTQQ8O8aLu1yOmnnx7mS5YsKT0T0P1WrFjR0yN0q8GDB4f5ySefHOYTJ04M87Fjx5b6uTfccEOYr1u3rtRzejMnSgBIKEoASChKAEgoSgBIKEoASPS5rde77747zCuVSpgXbbG22nZr//7x35k6Ozu7eRKojx122KGhzz/kkEPCvOhdM2bMmDDffffdw3zLLbcM8wkTJoR50f/Dv//+e5gvX748zP/8888wHzAgrou33347zPsSJ0oASChKAEgoSgBIKEoASChKAEg07dbrqaeeGuZtbW1hXq1Ww3zBggX1GqmpFW23Fv25vffeew2cBv5b0fZm0Xf0gQceCPNp06bVZZ4RI0aEedHW619//RXmv/32W5h/9NFHYT5r1qwwX7lyZZgXbfB3dHSE+erVq8N80KBBYb5q1aow70ucKAEgoSgBIKEoASChKAEgoSgBING0W69FG1hF9yN+//33Yf7EE0/UbabeZODAgWF+3XXXlXrOK6+8EuZXX3112ZGgS6ZMmRLmX3/9dZgfc8wxjRyn3zfffBPmzz77bJh//PHHYb5s2bJ6jVTKhRdeGOZDhw4N8y+//LKR4/RqTpQAkFCUAJBQlACQUJQAkFCUAJBo2q3Xsop+a/d3333XzZPUV9F2a3t7e5hfddVVYV50v+Ndd90V5hs3bqxhOmi82267radHaEonnnhiqc/PnTu3QZP0fk6UAJBQlACQUJQAkFCUAJBQlACQaJmt1wULFvT0CF3S1tYW5kVbrGeddVaYz58/P8zHjx+/WXMBrWHevHk9PUKPcaIEgISiBICEogSAhKIEgISiBIBE0269ViqVUvkZZ5wR5pdddlm9RqqLyy+/PMyvueaaMB8yZEiYz549O8wnTZq0eYMBtCgnSgBIKEoASChKAEgoSgBIKEoASDTt1mu1Wi2VDxs2LMzvueeeMJ81a1aY//jjj2F+1FFHhfm5554b5occckiY77777mH+zTffhPmiRYvCfMaMGWEOkCn6lwP7779/mC9btqyR4/QKTpQAkFCUAJBQlACQUJQAkFCUAJBo2q3XsrbYYoswnzJlSpiPHz8+zNevXx/mw4cP37zB/mHp0qVhvnjx4jC/9tpr6/JzAfr1K/6XA/37t+65qnX/ywGgBooSABKKEgASihIAEooSABJNu/X61ltvhfmKFSvCfNSoUaWeX3Q37C677FLqOUV3w86ZMyfML7vsslLPB+gORx99dJg/8sgj3TtID3CiBICEogSAhKIEgISiBICEogSARNNuva5evTrMx40bF+YXXXRRmLe3t9dlnunTp4f5/fffH+aff/55XX4uQD1VKpWeHqHXcaIEgISiBICEogSAhKIEgISiBIBEpVr066z/+UGbUPRiNX6N6QHeHb3T5MmTw3zWrFlh/tBDD4V50b8oaBa1vDucKAEgoSgBIKEoASChKAEgoSgBIGHrlT7B1mvv5d1Bb2brFQC6SFECQEJRAkBCUQJAQlECQEJRAkBCUQJAQlECQEJRAkBCUQJAQlECQEJRAkBCUQJAQlECQEJRAkBCUQJAQlECQKJS9avhAaCQEyUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkBtT6wUql0sg5oMuq1WpPj0DAu4PerJb3hhMlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkFCUAJBQlACQUJQAkBvT0ADROe3t7mF9//fVh3r9//Pem448/PsyXLFmyWXMBNBMnSgBIKEoASChKAEgoSgBIKEoASNh67QMmT54c5lOnTg3zzs7OUs+vVqtlRwLoM5woASChKAEgoSgBIKEoASChKAEgYeu1D9hzzz3DfKutturmSYCyjjzyyDCfOHFimI8ePTrMDzrooFI/98orrwzzNWvWhPmxxx4b5o8//niYL1++vNQ8vZkTJQAkFCUAJBQlACQUJQAkFCUAJCrVGi/yrFQqjZ6F/2HMmDFhPmfOnDAfMmRImK9atSrMTz311DDv6OgI8z/++CPMe4o7aXsn745/O+uss8J8+vTpYb7TTjuFedGf56uvvhrmQ4cODfMDDzwwzIsU/dynnnoqzM8+++xSz+8ptbw3nCgBIKEoASChKAEgoSgBIKEoASDhrtdeqOhOxYcffjjMi7Zbi9xxxx1h/vXXX5d6DrSyAQPi1+fhhx8e5g899FCYb7311mH+2muvhfkNN9wQ5m+88UaYDxw4MMyffPLJMB87dmyYF1m5cmWpzzcjJ0oASChKAEgoSgBIKEoASChKAEjYeu2FzjvvvDDfbbfdSj2n6O7Hxx57rOxIwD9MnDgxzGfOnFnqOS+99FKYF90Nu379+lLPL3pO2e3W1atXh/mjjz5a6jnNyIkSABKKEgASihIAEooSABKKEgASlWqNvxbebymvr6LfXt6vX79+HR0dYd7Z2Rnm69atC/MzzzwzzBcvXpwP16Rq/CrTzZr93VF0t+q0adPCvOh7OGPGjDBvb28P87LbrUU+/vjjMB8+fHip54wfPz7M58+fX3qm3qSW94YTJQAkFCUAJBQlACQUJQAkFCUAJNz12mB77bVXmM+dO7duP+Pee+8N87663QqNcO2114Z50Xbrpk2bwnzRokVhPnXq1DD//fffa5ju/2211VZhXnR36x577BHmRdvIN954Y5g3+3ZrVzhRAkBCUQJAQlECQEJRAkBCUQJAwtZrg5188slhPmLEiNLPevnll8N8+vTppZ8FrWq77bYL8ylTpoR50V2gRdutZ5xxxuaM9V/222+/MJ89e3aYjxw5stTzn3766TC//fbbSz2nFThRAkBCUQJAQlECQEJRAkBCUQJAolKt8dfCN/tvKW+0ok23Rx55JMy32WabwmctXbo0zM8888ww7+joSGdrFTV+lelmve3dsfPOO4f5mjVrSj1nn332CfM//vgjzM8///wwP+2008L84IMPDvNtt902zIu+/0X5uHHjwnzhwoVh3lfV8t5wogSAhKIEgISiBICEogSAhKIEgIS7Xkvaa6+9wnzu3Ll1+xlffvllmNtuha7btGlTmK9duzbMhw4dGuZfffVVmNdr+7poC3f9+vVhvuuuu4b5Dz/8EOattt3aFU6UAJBQlACQUJQAkFCUAJBQlACQsPVa0tSpU8O8s7Ozbj/j1ltvrduzgP+0bt26MC+6r/m5554L8x122CHMv/jiizCfP39+mBfdB/3TTz+F+Zw5c8K8aOu16PPUzokSABKKEgASihIAEooSABKKEgAStl4LtLW1hfnYsWPr8vyiDbh+/fr1++STT+ryM4DaLV++PMyL7npttOOOOy7MR48eHeZFm/dFd0dTOydKAEgoSgBIKEoASChKAEgoSgBI2Hot8OKLL4b59ttvX+o5y5YtC/PJkyeXHQloIYMGDQrzou3WarUa5u567TonSgBIKEoASChKAEgoSgBIKEoASNh6LbDjjjuGedHGWZEZM2aE+caNG0vPBLSORYsW9fQI/M2JEgASihIAEooSABKKEgASihIAEi2/9frwww+Hef/+9fk7xNKlS+vyHKC1nHTSST09An9zogSAhKIEgISiBICEogSAhKIEgETLbL22tbWF+ZgxY8K86E7XTZs2hfl9990X5h0dHf97OIB/2GeffXp6BP7mRAkACUUJAAlFCQAJRQkACUUJAImW2XrdbrvtwnzYsGGlnvPtt9+G+ZVXXll2JIBCr7/+epgX3UNdtKlP1zlRAkBCUQJAQlECQEJRAkBCUQJAomW2XgGayQcffBDmn332WZgX3Q277777hvnatWs3b7AW5EQJAAlFCQAJRQkACUUJAAlFCQCJltl6XbVqVZgvXbo0zI899thGjgOwWW6++eYwnzlzZpjfdNNNYX7ppZeG+UcffbR5g/VhTpQAkFCUAJBQlACQUJQAkFCUAJCoVKvVak0frFQaPQt0SY1fZbqZd0d9DR48OMyffPLJMB8zZkyYP/PMM2F+/vnnh/mvv/5aw3TNp5b3hhMlACQUJQAkFCUAJBQlACQUJQAkbL3SZ9h67Z28O7pH0TZs0V2vF198cZiPGDEizPvqHbC2XgGgixQlACQUJQAkFCUAJBQlACRsvdJn2Hrtnbw76M1svQJAFylKAEgoSgBIKEoASChKAEjUvPUKAK3IiRIAEooSABKKEgASihIAEooSABKKEgASihIAEooSABKKEgAS/wevhV4XL8ztagAAAABJRU5ErkJggg==\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# switch to False to use CPU\n",
        "use_cuda = True\n",
        "\n",
        "use_cuda = use_cuda and torch.cuda.is_available()\n",
        "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
        "torch.manual_seed(1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f8bCozymgXbu",
        "outputId": "221a59bb-81e6-4a14-f518-79529613f1e5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7fa649fe9250>"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Creating three tasks, task_1, task_2, and task_3, using the original MNIST dataset and permuted versions of the dataset.\n",
        "def permute_mnist(mnist, seed):\n",
        "    \"\"\" Given the training set, permute pixels of each img the same way. \"\"\"\n",
        "    np.random.seed(seed)\n",
        "    print(\"starting permutation...\")\n",
        "    h = w = 28\n",
        "    perm_inds = list(range(h*w))\n",
        "    np.random.shuffle(perm_inds)\n",
        "    # print(perm_inds)\n",
        "    perm_mnist = []\n",
        "    for dataset in mnist:\n",
        "        num_img = dataset.shape[0]\n",
        "        flat_set = dataset.reshape(num_img, w * h)\n",
        "        perm_mnist.append(flat_set[:, perm_inds].reshape(num_img, 1, w, h))\n",
        "    print(\"done.\")\n",
        "    return perm_mnist\n",
        "\n",
        "\n",
        "# task 1\n",
        "task_1 = [(x_train, t_train), (x_test, t_test)]\n",
        "\n",
        "# task 2\n",
        "x_train2, x_test2 = permute_mnist([x_train, x_test], 1)\n",
        "task_2 = [(x_train2, t_train), (x_test2, t_test)]\n",
        "\n",
        "# task 3\n",
        "x_train3, x_test3 = permute_mnist([x_train, x_test], 2)\n",
        "task_3 = [(x_train3, t_train), (x_test3, t_test)]\n",
        "\n",
        "# task list\n",
        "tasks = [task_1, task_2, task_3]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wyUXUnWRgacJ",
        "outputId": "7efd246e-d57f-40ad-ea22-07f6aec8ca7e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "starting permutation...\n",
            "done.\n",
            "starting permutation...\n",
            "done.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class Net(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
        "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
        "        self.conv2_drop = nn.Dropout2d()\n",
        "        self.fc1 = nn.Linear(320, 50)\n",
        "        self.fc2 = nn.Linear(50, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
        "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
        "        x = x.view(-1, 320)\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = F.dropout(x, training=self.training)\n",
        "        x = self.fc2(x)\n",
        "        return x"
      ],
      "metadata": {
        "id": "bXk-xclcgdEB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train_EWC(model, device, x_train, t_train, optimizer, epoch, ewc_lambda, fisher):\n",
        "    model.train()\n",
        "\n",
        "    for start in range(0, len(t_train)-1, 256):\n",
        "        end = start + 256\n",
        "        x, y = torch.from_numpy(x_train[start:end]), torch.from_numpy(t_train[start:end]).long()\n",
        "        x, y = x.to(device), y.to(device)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        output = model(x)\n",
        "        loss = F.cross_entropy(output, y)\n",
        "\n",
        "        loss.backward()  # Compute gradients\n",
        "\n",
        "        # Calculate EWC regularization term\n",
        "        ewc_regularization = 0.0\n",
        "        for name, param in model.named_parameters():\n",
        "            if name in fisher:\n",
        "                fisher[name] += torch.square(param.grad)\n",
        "                ewc_regularization += (ewc_lambda * fisher[name] * (param - param.detach()) ** 2).sum()\n",
        "\n",
        "        # Add EWC regularization term to the loss\n",
        "        loss += ewc_regularization\n",
        "\n",
        "        optimizer.step()\n",
        "\n",
        "    print('Train Epoch: {} \\tLoss: {:.6f}'.format(epoch, loss.item()))\n",
        "\n",
        "\n",
        "\n",
        "def test_EWC(model, device, x_test, t_test):\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    with torch.no_grad():\n",
        "        for start in range(0, len(t_test)-1, 256):\n",
        "            end = start + 256\n",
        "            x, y = torch.from_numpy(x_test[start:end]), torch.from_numpy(t_test[start:end]).long()\n",
        "            x, y = x.to(device), y.to(device)\n",
        "            output = model(x)\n",
        "            test_loss += F.cross_entropy(output, y).item()  # sum up batch loss\n",
        "            pred = output.max(1, keepdim=True)[1]  # get the index of the max logit\n",
        "            correct += pred.eq(y.view_as(pred)).sum().item()\n",
        "\n",
        "    test_loss /= len(t_test)\n",
        "    print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
        "        test_loss, correct, len(t_test),\n",
        "        100. * correct / len(t_test)))\n",
        "    return 100. * correct / len(t_test)"
      ],
      "metadata": {
        "id": "GB129ayAgf45"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = Net().to(device)\n",
        "optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)"
      ],
      "metadata": {
        "id": "pkq3yu45gkgQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# shuffle multiple lists (datasets) simultaneously\n",
        "def shuffle_in_unison(dataset, seed, in_place=False):\n",
        "    \"\"\" Shuffle two (or more) list in unison. \"\"\"\n",
        "\n",
        "    np.random.seed(seed)\n",
        "    rng_state = np.random.get_state()\n",
        "    new_dataset = []\n",
        "    for x in dataset:\n",
        "        if in_place:\n",
        "            np.random.shuffle(x)\n",
        "        else:\n",
        "            new_dataset.append(np.random.permutation(x))\n",
        "        np.random.set_state(rng_state)\n",
        "\n",
        "    if not in_place:\n",
        "        return new_dataset"
      ],
      "metadata": {
        "id": "n7Sn0V2zgmvt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "offline_accs = []\n",
        "\n",
        "# Store the Fisher information matrix for each parameter\n",
        "fisher = {}\n",
        "for name, param in model.named_parameters():\n",
        "    fisher[name] = torch.zeros_like(param)\n",
        "\n",
        "ewc_lambda = 0.01  # EWC regularization factor\n",
        "\n",
        "print(\"Training on all tasks together...\")\n",
        "avg_acc = 0\n",
        "(x_train, t_train), _ = tasks[0]\n",
        "\n",
        "for i in range(1, len(tasks)):\n",
        "    (past_x_train, past_t_train), _ = tasks[i]\n",
        "    x_train = np.concatenate((x_train, past_x_train))\n",
        "    t_train = np.concatenate((t_train, past_t_train))\n",
        "\n",
        "x_train, t_train = shuffle_in_unison([x_train, t_train], 0)\n",
        "\n",
        "for epoch in range(1, 4):\n",
        "    train_EWC(model, device, x_train, t_train, optimizer, epoch, ewc_lambda, fisher)\n",
        "\n",
        "for id_test, task in enumerate(tasks):\n",
        "    print(\"Testing on task: \", id_test)\n",
        "    _, (x_test, t_test) = task\n",
        "    acc = test_EWC(model, device, x_test, t_test)\n",
        "    avg_acc = avg_acc + acc\n",
        "\n",
        "print(\"Avg acc: \", avg_acc / 3)\n",
        "for i in range(len(tasks)):\n",
        "    offline_accs.append(avg_acc/3)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3JF9IksUgrc3",
        "outputId": "32e3066d-23ec-4246-dd17-0a85a778053a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training on all tasks together...\n",
            "Train Epoch: 1 \tLoss: 1.346628\n",
            "Train Epoch: 2 \tLoss: 0.810519\n",
            "Train Epoch: 3 \tLoss: 0.831743\n",
            "Testing on task:  0\n",
            "Test set: Average loss: 0.0007, Accuracy: 9440/10000 (94%)\n",
            "\n",
            "Testing on task:  1\n",
            "Test set: Average loss: 0.0017, Accuracy: 8792/10000 (88%)\n",
            "\n",
            "Testing on task:  2\n",
            "Test set: Average loss: 0.0018, Accuracy: 8729/10000 (87%)\n",
            "\n",
            "Avg acc:  89.87\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "JZIq61NqguST"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}