Elastic_Weight_Consolidation_(EWC).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyNNl81TdHGC/n6JvwYwbZ5O",
      "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/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": "Qg22loB7zk9z"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from transformers import AutoModelForSequenceClassification, AutoTokenizer\n",
        "\n",
        "class EWC:\n",
        "    def __init__(self, model, importance):\n",
        "        self.model = model\n",
        "        self.importance = importance\n",
        "        self.fisher_information = {}\n",
        "        self.params_old = {}\n",
        "\n",
        "    def compute_fisher_information(self, data_loader, criterion):\n",
        "        self.model.eval()\n",
        "\n",
        "        # Initialize Fisher information\n",
        "        for name, param in self.model.named_parameters():\n",
        "            self.fisher_information[name] = torch.zeros_like(param)\n",
        "            self.params_old[name] = param.clone().detach()\n",
        "\n",
        "        # Compute Fisher Information for each parameter\n",
        "        for inputs, labels in data_loader:\n",
        "            self.model.zero_grad()\n",
        "            outputs = self.model(inputs)\n",
        "            loss = criterion(outputs.logits, labels)\n",
        "            loss.backward()\n",
        "\n",
        "            for name, param in self.model.named_parameters():\n",
        "                self.fisher_information[name] += param.grad.data.pow(2) / len(data_loader)\n",
        "\n",
        "        self.model.train()\n",
        "\n",
        "    def penalty(self):\n",
        "        penalty = 0\n",
        "        for name, param in self.model.named_parameters():\n",
        "            fisher_val = self.fisher_information[name]\n",
        "            param_old = self.params_old[name]\n",
        "            penalty += (self.importance * fisher_val * (param - param_old).pow(2)).sum()\n",
        "        return penalty\n",
        "\n",
        "# Example usage\n",
        "# Create a dummy dataset\n",
        "input_data = torch.randint(0, 1000, (100, 10))  # 100 samples, sequence length 10\n",
        "labels = torch.randint(0, 2, (100,))  # 100 binary labels\n",
        "\n",
        "dataset = TensorDataset(input_data, labels)\n",
        "data_loader = DataLoader(dataset, batch_size=10)\n",
        "\n",
        "# Initialize model and tokenizer\n",
        "model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-uncased\", num_labels=2)\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)\n",
        "\n",
        "ewc = EWC(model, importance=1.0)\n",
        "\n",
        "# Compute Fisher Information\n",
        "ewc.compute_fisher_information(data_loader, criterion)\n",
        "\n",
        "# Use the penalty in the training loop\n",
        "for inputs, labels in data_loader:\n",
        "    model.zero_grad()\n",
        "    outputs = model(inputs)[\"logits\"]\n",
        "    loss = criterion(outputs, labels) + ewc.penalty()  # Add EWC penalty to the loss\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    print(f\"Loss: {loss.item()}\")"
      ]
    }
  ]
}