Adversarial_Training_with_PGD_(Projected_Gradient_Descent).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyMkKriKgUDpvDlFsjMG+4l8",
      "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_with_PGD_(Projected_Gradient_Descent).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": "tYJXFrqBTZ1P"
      },
      "outputs": [],
      "source": [
        "pip install torchattacks"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "from torchattacks import PGD\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "# Define a simple CNN model\n",
        "class SimpleCNN(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(SimpleCNN, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)\n",
        "        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)\n",
        "        self.fc1 = nn.Linear(64 * 12 * 12, 128)\n",
        "        self.fc2 = nn.Linear(128, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(self.conv1(x))\n",
        "        x = F.relu(self.conv2(x))\n",
        "        x = F.max_pool2d(x, 2)\n",
        "        x = x.view(-1, 64 * 12 * 12)\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = self.fc2(x)\n",
        "        return x\n",
        "\n",
        "# Load dataset\n",
        "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])\n",
        "train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
        "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
        "\n",
        "# Instantiate the model, loss function, and optimizer\n",
        "model = SimpleCNN()\n",
        "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "# Define the PGD attack\n",
        "pgd_attack = PGD(model, eps=0.3, alpha=2/255, steps=40)\n",
        "\n",
        "# Number of epochs\n",
        "epochs = 5\n",
        "\n",
        "# Adversarial training setup without CUDA\n",
        "for epoch in range(epochs):\n",
        "    model.train()\n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        # Move data and target to CPU\n",
        "        data, target = data.to('cpu'), target.to('cpu')\n",
        "\n",
        "        # Generate adversarial examples\n",
        "        adv_data = pgd_attack(data, target)\n",
        "\n",
        "        # Forward pass and loss calculation\n",
        "        optimizer.zero_grad()\n",
        "        output = model(adv_data)\n",
        "        loss = criterion(output, target)\n",
        "\n",
        "        # Backward pass and optimization\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        if batch_idx % 100 == 0:\n",
        "            print(f'Epoch [{epoch + 1}/{epochs}], Batch [{batch_idx}/{len(train_loader)}], Loss: {loss.item():.4f}')\n",
        "\n",
        "print(\"Adversarial training completed successfully.\")"
      ],
      "metadata": {
        "id": "RtJIswi4TmC5"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}