MoE2.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b70a7041-53d9-4305-bdb0-d22173b50a95",
      "metadata": {
        "id": "b70a7041-53d9-4305-bdb0-d22173b50a95"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "208619ce-0270-4e6d-b30f-90107432e0ba",
      "metadata": {
        "id": "208619ce-0270-4e6d-b30f-90107432e0ba"
      },
      "outputs": [],
      "source": [
        "# Define the expert model\n",
        "class Expert(nn.Module):\n",
        "    def __init__(self, input_dim, hidden_dim, output_dim):\n",
        "        super(Expert, self).__init__()\n",
        "        self.layer1 = nn.Linear(input_dim, hidden_dim)\n",
        "        self.layer2 = nn.Linear(hidden_dim, output_dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = torch.relu(self.layer1(x))\n",
        "        return torch.softmax(self.layer2(x), dim=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "078fb12a-60d5-4e5f-8d72-a49d01aa9291",
      "metadata": {
        "id": "078fb12a-60d5-4e5f-8d72-a49d01aa9291"
      },
      "outputs": [],
      "source": [
        "# Define the gating model\n",
        "class Gating(nn.Module):\n",
        "    def __init__(self, input_dim,\n",
        "                 num_experts, dropout_rate=0.1):\n",
        "        super(Gating, self).__init__()\n",
        "\n",
        "        # Layers\n",
        "        self.layer1 = nn.Linear(input_dim, 128)\n",
        "        self.dropout1 = nn.Dropout(dropout_rate)\n",
        "\n",
        "        self.layer2 = nn.Linear(128, 256)\n",
        "        self.leaky_relu1 = nn.LeakyReLU()\n",
        "        self.dropout2 = nn.Dropout(dropout_rate)\n",
        "\n",
        "        self.layer3 = nn.Linear(256, 128)\n",
        "        self.leaky_relu2 = nn.LeakyReLU()\n",
        "        self.dropout3 = nn.Dropout(dropout_rate)\n",
        "\n",
        "        self.layer4 = nn.Linear(128, num_experts)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = torch.relu(self.layer1(x))\n",
        "        x = self.dropout1(x)\n",
        "\n",
        "        x = self.layer2(x)\n",
        "        x = self.leaky_relu1(x)\n",
        "        x = self.dropout2(x)\n",
        "\n",
        "        x = self.layer3(x)\n",
        "        x = self.leaky_relu2(x)\n",
        "        x = self.dropout3(x)\n",
        "\n",
        "        return torch.softmax(self.layer4(x), dim=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ba02da85",
      "metadata": {
        "id": "ba02da85"
      },
      "outputs": [],
      "source": [
        "class MoE(nn.Module):\n",
        "    def __init__(self, trained_experts):\n",
        "        super(MoE, self).__init__()\n",
        "        self.experts = nn.ModuleList(trained_experts)\n",
        "\n",
        "        # Freezing the experts to ensure that they are not\n",
        "        # learning when MoE is training.\n",
        "        # Ideally, one can free them before sending the\n",
        "        # experts to the MoE; in that case the following three\n",
        "        # lines can be commented out.\n",
        "        for expert in self.experts:\n",
        "            for param in expert.parameters():\n",
        "                param.requires_grad = False\n",
        "\n",
        "        num_experts = len(trained_experts)\n",
        "        # Assuming all experts have the same input dimension\n",
        "        input_dim = trained_experts[0].layer1.in_features\n",
        "        self.gating = Gating(input_dim, num_experts)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Get the weights from the gating network\n",
        "        weights = self.gating(x)\n",
        "\n",
        "        # Calculate the expert outputs\n",
        "        outputs = torch.stack(\n",
        "            [expert(x) for expert in self.experts], dim=2)\n",
        "\n",
        "        # Adjust the weights tensor shape to match the expert outputs\n",
        "        weights = weights.unsqueeze(1).expand_as(outputs)\n",
        "\n",
        "        # Multiply the expert outputs with the weights and\n",
        "        # sum along the third dimension\n",
        "        return torch.sum(outputs * weights, dim=2)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9c84da01-bdb4-4e7e-af89-f2ce617f4da0",
      "metadata": {
        "id": "9c84da01-bdb4-4e7e-af89-f2ce617f4da0",
        "outputId": "8b26d1a6-e0e7-4128-b041-1486d1b0c61c",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([1999, 4]) \n",
            " torch.Size([500, 4]) \n",
            " torch.Size([1658, 4]) \n",
            " torch.Size([1658, 4]) \n",
            " torch.Size([1658, 4])\n"
          ]
        }
      ],
      "source": [
        "# Generate the dataset\n",
        "num_samples = 5000\n",
        "input_dim = 4\n",
        "hidden_dim = 32\n",
        "\n",
        "# Generate equal numbers of labels 0, 1, and 2\n",
        "y_data = torch.cat([\n",
        "    torch.zeros(num_samples // 3),\n",
        "    torch.ones(num_samples // 3),\n",
        "    torch.full((num_samples - 2 * (num_samples // 3),), 2)  # Filling the remaining to ensure exact num_samples\n",
        "]).long()\n",
        "\n",
        "# Biasing the data based on the labels\n",
        "x_data = torch.randn(num_samples, input_dim)\n",
        "\n",
        "for i in range(num_samples):\n",
        "    if y_data[i] == 0:\n",
        "        x_data[i, 0] += 1  # Making x[0] more positive\n",
        "    elif y_data[i] == 1:\n",
        "        x_data[i, 1] -= 1  # Making x[1] more negative\n",
        "    elif y_data[i] == 2:\n",
        "        x_data[i, 0] -= 1  # Making x[0] more negative\n",
        "\n",
        "# Shuffle the data to randomize the order\n",
        "indices = torch.randperm(num_samples)\n",
        "x_data = x_data[indices]\n",
        "y_data = y_data[indices]\n",
        "\n",
        "# Verify the label distribution\n",
        "y_data.bincount()\n",
        "\n",
        "# Shuffle the data to ensure x_data and y_data remain aligned\n",
        "shuffled_indices = torch.randperm(num_samples)\n",
        "x_data = x_data[shuffled_indices]\n",
        "y_data = y_data[shuffled_indices]\n",
        "\n",
        "# Splitting data for training individual experts\n",
        "# Use the first half samples for training individual experts\n",
        "x_train_experts = x_data[:int(num_samples/2)]\n",
        "y_train_experts = y_data[:int(num_samples/2)]\n",
        "\n",
        "mask_expert1 = (y_train_experts == 0) | (y_train_experts == 1)\n",
        "mask_expert2 = (y_train_experts == 1) | (y_train_experts == 2)\n",
        "mask_expert3 = (y_train_experts == 0) | (y_train_experts == 2)\n",
        "\n",
        "# Select an almost equal number of samples for each expert\n",
        "num_samples_per_expert = \\\n",
        "min(mask_expert1.sum(), mask_expert2.sum(), mask_expert3.sum())\n",
        "\n",
        "x_expert1 = x_train_experts[mask_expert1][:num_samples_per_expert]\n",
        "y_expert1 = y_train_experts[mask_expert1][:num_samples_per_expert]\n",
        "\n",
        "x_expert2 = x_train_experts[mask_expert2][:num_samples_per_expert]\n",
        "y_expert2 = y_train_experts[mask_expert2][:num_samples_per_expert]\n",
        "\n",
        "x_expert3 = x_train_experts[mask_expert3][:num_samples_per_expert]\n",
        "y_expert3 = y_train_experts[mask_expert3][:num_samples_per_expert]\n",
        "\n",
        "# Splitting the next half samples for training MoE model and for testing\n",
        "x_remaining = x_data[int(num_samples/2)+1:]\n",
        "y_remaining = y_data[int(num_samples/2)+1:]\n",
        "\n",
        "split = int(0.8 * len(x_remaining))\n",
        "x_train_moe = x_remaining[:split]\n",
        "y_train_moe = y_remaining[:split]\n",
        "\n",
        "x_test = x_remaining[split:]\n",
        "y_test = y_remaining[split:]\n",
        "\n",
        "print(x_train_moe.shape,\"\\n\", x_test.shape,\"\\n\",\n",
        "      x_expert1.shape,\"\\n\",\n",
        "      x_expert2.shape,\"\\n\", x_expert3.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b34f8f45-cc9b-4570-b7a5-f8d767b88ed6",
      "metadata": {
        "id": "b34f8f45-cc9b-4570-b7a5-f8d767b88ed6"
      },
      "outputs": [],
      "source": [
        "# Define hidden dimension\n",
        "output_dim = 3\n",
        "hidden_dim = 32\n",
        "\n",
        "epochs = 500\n",
        "learning_rate = 0.001\n",
        "\n",
        "\n",
        "# Instantiate the experts\n",
        "expert1 = Expert(input_dim, hidden_dim, output_dim)\n",
        "expert2 = Expert(input_dim, hidden_dim, output_dim)\n",
        "expert3 = Expert(input_dim, hidden_dim, output_dim)\n",
        "\n",
        "# Set up loss\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "# Optimizers for experts\n",
        "optimizer_expert1 = optim.Adam(expert1.parameters(), lr=learning_rate)\n",
        "optimizer_expert2 = optim.Adam(expert2.parameters(), lr=learning_rate)\n",
        "optimizer_expert3 = optim.Adam(expert3.parameters(), lr=learning_rate)\n",
        "\n",
        "# Training loop for expert 1\n",
        "for epoch in range(epochs):\n",
        "    optimizer_expert1.zero_grad()\n",
        "    outputs_expert1 = expert1(x_expert1)\n",
        "    loss_expert1 = criterion(outputs_expert1, y_expert1)\n",
        "    loss_expert1.backward()\n",
        "    optimizer_expert1.step()\n",
        "\n",
        "# Training loop for expert 2\n",
        "for epoch in range(epochs):\n",
        "    optimizer_expert2.zero_grad()\n",
        "    outputs_expert2 = expert2(x_expert2)\n",
        "    loss_expert2 = criterion(outputs_expert2, y_expert2)\n",
        "    loss_expert2.backward()\n",
        "    optimizer_expert2.step()\n",
        "\n",
        "# Training loop for expert 3\n",
        "for epoch in range(epochs):\n",
        "    optimizer_expert3.zero_grad()\n",
        "    outputs_expert3 = expert3(x_expert3)\n",
        "    loss_expert3 = criterion(outputs_expert3, y_expert3)\n",
        "    loss_expert3.backward()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "208f3673-a3aa-4e1f-b55e-d53efa7d3ed2",
      "metadata": {
        "id": "208f3673-a3aa-4e1f-b55e-d53efa7d3ed2"
      },
      "outputs": [],
      "source": [
        "# Create the MoE model with the trained experts\n",
        "moe_model = MoE([expert1, expert2, expert3])\n",
        "\n",
        "# Train the MoE model\n",
        "optimizer_moe = optim.Adam(moe_model.parameters(), lr=learning_rate)\n",
        "for epoch in range(epochs):\n",
        "    optimizer_moe.zero_grad()\n",
        "    outputs_moe = moe_model(x_train_moe)\n",
        "    loss_moe = criterion(outputs_moe, y_train_moe)\n",
        "    loss_moe.backward()\n",
        "    optimizer_moe.step()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5ac480d5-2a8d-4b77-9223-f7ac0feedc30",
      "metadata": {
        "id": "5ac480d5-2a8d-4b77-9223-f7ac0feedc30"
      },
      "outputs": [],
      "source": [
        "# Evaluate all models\n",
        "def evaluate(model, x, y):\n",
        "    with torch.no_grad():\n",
        "        outputs = model(x)\n",
        "        _, predicted = torch.max(outputs, 1)\n",
        "        correct = (predicted == y).sum().item()\n",
        "        accuracy = correct / len(y)\n",
        "    return accuracy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3c643949-3445-42df-ad02-ff4f945fb601",
      "metadata": {
        "id": "3c643949-3445-42df-ad02-ff4f945fb601",
        "outputId": "94aeae3b-3338-4415-fd94-c13421c1f5f5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Expert 1 Accuracy: 0.466\n",
            "Expert 2 Accuracy: 0.496\n",
            "Expert 3 Accuracy: 0.378\n",
            "Mixture of Experts Accuracy: 0.614\n"
          ]
        }
      ],
      "source": [
        "accuracy_expert1 = evaluate(expert1, x_test, y_test)\n",
        "accuracy_expert2 = evaluate(expert2, x_test, y_test)\n",
        "accuracy_expert3 = evaluate(expert3, x_test, y_test)\n",
        "accuracy_moe = evaluate(moe_model, x_test, y_test)\n",
        "\n",
        "print(\"Expert 1 Accuracy:\", accuracy_expert1)\n",
        "print(\"Expert 2 Accuracy:\", accuracy_expert2)\n",
        "print(\"Expert 3 Accuracy:\", accuracy_expert3)\n",
        "print(\"Mixture of Experts Accuracy:\", accuracy_moe)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ec4c217a-edc6-493b-88fe-fdb187f6e202",
      "metadata": {
        "id": "ec4c217a-edc6-493b-88fe-fdb187f6e202"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "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.11.4"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}