Continual_Learning_with_Elastic_Weight_Consolidation_(EWC).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyO01JicDaaXJyFDKHvfnbld",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/OneFineStarstuff/OneFineStardust/blob/main/Continual_Learning_with_Elastic_Weight_Consolidation_(EWC).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-iVBDfoFhHX2"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "\n",
        "# Define a simple model\n",
        "class MyModel(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(MyModel, self).__init__()\n",
        "        self.fc1 = nn.Linear(784, 128)\n",
        "        self.fc2 = nn.Linear(128, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = torch.relu(self.fc1(x))\n",
        "        x = self.fc2(x)\n",
        "        return x\n",
        "\n",
        "class EWC:\n",
        "    def __init__(self, model, dataloader, importance):\n",
        "        self.model = model\n",
        "        self.dataloader = dataloader\n",
        "        self.importance = importance\n",
        "        self.params = {n: p for n, p in self.model.named_parameters() if p.requires_grad}\n",
        "        self._means = {}\n",
        "        self._fisher = {}\n",
        "        self._compute_fisher()\n",
        "\n",
        "    def _compute_fisher(self):\n",
        "        self.model.eval()\n",
        "        fisher = {n: torch.zeros_like(p) for n, p in self.params.items()}\n",
        "        for inputs, labels in self.dataloader:\n",
        "            self.model.zero_grad()\n",
        "            outputs = self.model(inputs)\n",
        "            loss = nn.CrossEntropyLoss()(outputs, labels)\n",
        "            loss.backward()\n",
        "            for n, p in self.params.items():\n",
        "                fisher[n] += p.grad.data.pow(2) / len(self.dataloader)\n",
        "        for n, p in fisher.items():\n",
        "            self._fisher[n] = fisher[n]\n",
        "            self._means[n] = self.params[n].data.clone()\n",
        "\n",
        "    def penalty(self, model):\n",
        "        loss = 0\n",
        "        for n, p in model.named_parameters():\n",
        "            if n in self._fisher:\n",
        "                _loss = self._fisher[n] * (p - self._means[n]).pow(2)\n",
        "                loss += _loss.sum()\n",
        "        return self.importance * loss\n",
        "\n",
        "    def update_model(self, new_data, optimizer):\n",
        "        self.model.train()\n",
        "        dataloader = DataLoader(TensorDataset(*new_data), batch_size=32, shuffle=True)\n",
        "        for inputs, labels in dataloader:\n",
        "            optimizer.zero_grad()\n",
        "            outputs = self.model(inputs)\n",
        "            loss = nn.CrossEntropyLoss()(outputs, labels)\n",
        "            ewc_loss = loss + self.penalty(self.model)\n",
        "            ewc_loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "# Initialize your model, dataloader, and optimizer\n",
        "model = MyModel()\n",
        "X_train = torch.randn(1000, 784)  # example training data\n",
        "y_train = torch.randint(0, 10, (1000,))  # example training labels\n",
        "dataloader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)\n",
        "optimizer = optim.SGD(model.parameters(), lr=0.01)\n",
        "\n",
        "# Initialize EWC\n",
        "ewc = EWC(model, dataloader, importance=1000)\n",
        "\n",
        "# Example new data stream for continual learning\n",
        "new_X1 = torch.randn(200, 784)  # example new data\n",
        "new_y1 = torch.randint(0, 10, (200,))  # example new labels\n",
        "data_stream = [(new_X1, new_y1)]\n",
        "\n",
        "# Applying EWC in the learning loop\n",
        "for new_data in data_stream:\n",
        "    ewc.update_model(new_data, optimizer)"
      ]
    }
  ]
}