Adversarial_Training_for_Enhanced_Security.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyMyDvdTk3iAwCrMlX3RIuvm",
      "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/Adversarial_Training_for_Enhanced_Security.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "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 the model (example: simple feedforward neural network)\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",
        "# Generate adversarial examples using FGSM (Fast Gradient Sign Method)\n",
        "def generate_adversarial_example(model, data, target, epsilon):\n",
        "    data.requires_grad = True\n",
        "    output = model(data)\n",
        "    loss = nn.CrossEntropyLoss()(output, target)\n",
        "    model.zero_grad()\n",
        "    loss.backward()\n",
        "    data_grad = data.grad.data\n",
        "    adv_data = data + epsilon * data_grad.sign()\n",
        "    adv_data = torch.clamp(adv_data, 0, 1)  # Keep data within valid range\n",
        "    return adv_data\n",
        "\n",
        "# Define the training process\n",
        "def adversarial_training(model, data_loader, optimizer, epsilon, device):\n",
        "    model.train()\n",
        "    for data, target in data_loader:\n",
        "        data, target = data.to(device), target.to(device)\n",
        "\n",
        "        # Generate adversarial examples\n",
        "        adv_data = generate_adversarial_example(model, data, target, epsilon)\n",
        "\n",
        "        # Training on adversarial examples\n",
        "        optimizer.zero_grad()\n",
        "        output = model(adv_data)\n",
        "        loss = nn.CrossEntropyLoss()(output, target)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "# Main script to set up model, data, and training loop\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "model = MyModel().to(device)\n",
        "optimizer = optim.SGD(model.parameters(), lr=0.01)\n",
        "\n",
        "# Example data (replace with your actual dataset)\n",
        "X_train = torch.randn(1000, 784)  # Example data\n",
        "y_train = torch.randint(0, 10, (1000,))  # Example labels\n",
        "train_dataset = TensorDataset(X_train, y_train)\n",
        "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
        "\n",
        "# Set the adversarial attack strength\n",
        "epsilon = 0.1\n",
        "\n",
        "# Train the model with adversarial examples\n",
        "adversarial_training(model, train_loader, optimizer, epsilon, device)"
      ],
      "metadata": {
        "id": "cjosvj2eXH3b"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}