{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### train a binary classifier for CIFAKE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "d:\\Anaconda\\envs\\ai\\Lib\\site-packages\\torchvision\\models\\_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "d:\\Anaconda\\envs\\ai\\Lib\\site-packages\\torchvision\\models\\_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "Train Loss: 0.4615 Acc: 0.7786\n",
      "Epoch 2/10\n",
      "Train Loss: 0.3657 Acc: 0.8369\n",
      "Epoch 3/10\n",
      "Train Loss: 0.3367 Acc: 0.8522\n",
      "Epoch 4/10\n",
      "Train Loss: 0.3150 Acc: 0.8626\n",
      "Epoch 5/10\n",
      "Train Loss: 0.3046 Acc: 0.8680\n",
      "Epoch 6/10\n",
      "Train Loss: 0.2917 Acc: 0.8745\n",
      "Epoch 7/10\n",
      "Train Loss: 0.2805 Acc: 0.8809\n",
      "Epoch 8/10\n",
      "Train Loss: 0.2760 Acc: 0.8824\n",
      "Epoch 9/10\n",
      "Train Loss: 0.2714 Acc: 0.8843\n",
      "Epoch 10/10\n",
      "Train Loss: 0.2653 Acc: 0.8868\n",
      "Model saved to vgg_model.pth\n",
      "Test Accuracy: 0.8859\n",
      "Final test accuracy: 0.8859\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "from torchvision import datasets, models, transforms\n",
    "\n",
    "def train_and_save_vgg_model(data_dir, model_path, num_epochs=30, batch_size=32, lr=0.00001):\n",
    "    data_transform = transforms.Compose([\n",
    "        transforms.RandomResizedCrop(32),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.RandomRotation(15),\n",
    "        transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])\n",
    "    ])\n",
    "\n",
    "    full_dataset = datasets.ImageFolder(data_dir, transform=data_transform)\n",
    "\n",
    "    train_size = int(0.8 * len(full_dataset))\n",
    "    test_size = len(full_dataset) - train_size\n",
    "    train_dataset, test_dataset = random_split(full_dataset, [train_size, test_size])\n",
    "\n",
    "    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)\n",
    "\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "    model = models.vgg16(pretrained=True)\n",
    "    for param in model.features.parameters(): \n",
    "        param.requires_grad = True\n",
    "    model.classifier[6] = nn.Linear(model.classifier[6].in_features, 2)\n",
    "    model = model.to(device)\n",
    "\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        print(f\"Epoch {epoch + 1}/{num_epochs}\")\n",
    "        model.train()\n",
    "\n",
    "        running_loss = 0.0\n",
    "        running_corrects = 0\n",
    "\n",
    "        for inputs, labels in train_loader:\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            running_loss += loss.item() * inputs.size(0)\n",
    "            running_corrects += torch.sum(preds == labels.data)\n",
    "\n",
    "        epoch_loss = running_loss / train_size\n",
    "        epoch_acc = running_corrects.double() / train_size\n",
    "\n",
    "        print(f\"Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}\")\n",
    "        scheduler.step()\n",
    "\n",
    "    torch.save(model.state_dict(), model_path)\n",
    "    print(f\"Model saved to {model_path}\")\n",
    "\n",
    "    model.eval()\n",
    "    test_corrects = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in test_loader:\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "\n",
    "            outputs = model(inputs)\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "            test_corrects += torch.sum(preds == labels.data)\n",
    "\n",
    "    test_acc = test_corrects.double() / test_size\n",
    "    print(f\"Test Accuracy: {test_acc:.4f}\")\n",
    "\n",
    "    return test_acc.item()\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    data_dir = \"CIFAKE/train\"  \n",
    "    model_path = \"vgg_model.pth\"\n",
    "    num_epochs = 10\n",
    "    batch_size = 32\n",
    "    lr = 0.00001\n",
    "\n",
    "    test_acc = train_and_save_vgg_model(data_dir, model_path, num_epochs, batch_size, lr)\n",
    "    print(f\"Final test accuracy: {test_acc:.4f}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}