{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://alirezasamar.com/blog/2023/03/fine-tuning-pre-trained-resnet-18-model-image-classification-pytorch/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\env\\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",
      "c:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\env\\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=ResNet18_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet18_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision.models as models\n",
    "\n",
    "# Load the pre-trained ResNet-18 model\n",
    "model = models.resnet18(pretrained=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Freeze all the pre-trained layers\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modify the last layer of the model\n",
    "num_classes = 13 # replace with the number of classes in your dataset\n",
    "model.fc = torch.nn.Linear(model.fc.in_features, num_classes)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.datasets import ImageFolder\n",
    "from torchvision.transforms import transforms\n",
    "\n",
    "# Define the transformations to apply to the images\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "# Load the train and validation datasets\n",
    "train_dataset = ImageFolder(r'D:\\color_classify\\dataset_split\\train', transform=transform)\n",
    "val_dataset = ImageFolder(r'D:\\color_classify\\dataset_split\\val', transform=transform)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the loss function and optimizer\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# Create data loaders for the train and validation datasets\n",
    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print (device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader, val_loader, criterion, optimizer, num_epochs):\n",
    "    # Train the model for the specified number of epochs\n",
    "    for epoch in range(num_epochs):\n",
    "        # Set the model to train mode\n",
    "        model.train()\n",
    "\n",
    "        # Initialize the running loss and accuracy\n",
    "        running_loss = 0.0\n",
    "        running_corrects = 0\n",
    "\n",
    "        # Iterate over the batches of the train loader\n",
    "        for inputs, labels in train_loader:\n",
    "            # Move the inputs and labels to the device\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            # Zero the optimizer gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # Forward pass\n",
    "            outputs = model(inputs)\n",
    "            _, preds = torch.max(outputs, 1)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            # Backward pass and optimizer step\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # Update the running loss and accuracy\n",
    "            running_loss += loss.item() * inputs.size(0)\n",
    "            running_corrects += torch.sum(preds == labels.data)\n",
    "\n",
    "        # Calculate the train loss and accuracy\n",
    "        train_loss = running_loss / len(train_dataset)\n",
    "        train_acc = running_corrects.double() / len(train_dataset)\n",
    "\n",
    "        # Set the model to evaluation mode\n",
    "        model.eval()\n",
    "\n",
    "        # Initialize the running loss and accuracy\n",
    "        running_loss = 0.0\n",
    "        running_corrects = 0\n",
    "\n",
    "        # Iterate over the batches of the validation loader\n",
    "        with torch.no_grad():\n",
    "            for inputs, labels in val_loader:\n",
    "                # Move the inputs and labels to the device\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "\n",
    "                # Forward pass\n",
    "                outputs = model(inputs)\n",
    "                _, preds = torch.max(outputs, 1)\n",
    "                loss = criterion(outputs, labels)\n",
    "\n",
    "                # Update the running loss and accuracy\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                running_corrects += torch.sum(preds == labels.data)\n",
    "\n",
    "        # Calculate the validation loss and accuracy\n",
    "        val_loss = running_loss / len(val_dataset)\n",
    "        val_acc = running_corrects.double() / len(val_dataset)\n",
    "\n",
    "        # Print the epoch results\n",
    "        print('Epoch [{}/{}], train loss: {:.4f}, train acc: {:.4f}, val loss: {:.4f}, val acc: {:.4f}'\n",
    "              .format(epoch+1, num_epochs, train_loss, train_acc, val_loss, val_acc))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "# Set the device\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model.to(device)\n",
    "print (device)\n",
    "# Fine-tune the last layer for a few epochs\n",
    "optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.01, momentum=0.9)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[10], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mval_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcriterion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_epochs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m5\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;66;03m# Unfreeze all the layers and fine-tune the entire network for a few more epochs\u001b[39;00m\n\u001b[0;32m      4\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m param \u001b[38;5;129;01min\u001b[39;00m model\u001b[38;5;241m.\u001b[39mparameters():\n",
      "Cell \u001b[1;32mIn[8], line 12\u001b[0m, in \u001b[0;36mtrain\u001b[1;34m(model, train_loader, val_loader, criterion, optimizer, num_epochs)\u001b[0m\n\u001b[0;32m      9\u001b[0m running_corrects \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[0;32m     11\u001b[0m \u001b[38;5;66;03m# Iterate over the batches of the train loader\u001b[39;00m\n\u001b[1;32m---> 12\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m inputs, labels \u001b[38;5;129;01min\u001b[39;00m train_loader:\n\u001b[0;32m     13\u001b[0m     \u001b[38;5;66;03m# Move the inputs and labels to the device\u001b[39;00m\n\u001b[0;32m     14\u001b[0m     inputs \u001b[38;5;241m=\u001b[39m inputs\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[0;32m     15\u001b[0m     labels \u001b[38;5;241m=\u001b[39m labels\u001b[38;5;241m.\u001b[39mto(device)\n",
      "File \u001b[1;32mc:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\env\\lib\\site-packages\\torch\\utils\\data\\dataloader.py:708\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    705\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    706\u001b[0m     \u001b[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[39;00m\n\u001b[0;32m    707\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()  \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[1;32m--> 708\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_next_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    709\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[0;32m    710\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[0;32m    711\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable\n\u001b[0;32m    712\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m    713\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called\n\u001b[0;32m    714\u001b[0m ):\n",
      "File \u001b[1;32mc:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\env\\lib\\site-packages\\torch\\utils\\data\\dataloader.py:764\u001b[0m, in \u001b[0;36m_SingleProcessDataLoaderIter._next_data\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    762\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_next_data\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[0;32m    763\u001b[0m     index \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_next_index()  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[1;32m--> 764\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_dataset_fetcher\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfetch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mindex\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[0;32m    765\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory:\n\u001b[0;32m    766\u001b[0m         data \u001b[38;5;241m=\u001b[39m _utils\u001b[38;5;241m.\u001b[39mpin_memory\u001b[38;5;241m.\u001b[39mpin_memory(data, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory_device)\n",
      "File \u001b[1;32mc:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\env\\lib\\site-packages\\torch\\utils\\data\\_utils\\fetch.py:52\u001b[0m, in \u001b[0;36m_MapDatasetFetcher.fetch\u001b[1;34m(self, possibly_batched_index)\u001b[0m\n\u001b[0;32m     50\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset\u001b[38;5;241m.\u001b[39m__getitems__(possibly_batched_index)\n\u001b[0;32m     51\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m---> 52\u001b[0m         data \u001b[38;5;241m=\u001b[39m [\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[idx] \u001b[38;5;28;01mfor\u001b[39;00m idx \u001b[38;5;129;01min\u001b[39;00m possibly_batched_index]\n\u001b[0;32m     53\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m     54\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[possibly_batched_index]\n",
      "File \u001b[1;32mc:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\env\\lib\\site-packages\\torch\\utils\\data\\_utils\\fetch.py:52\u001b[0m, in \u001b[0;36m<listcomp>\u001b[1;34m(.0)\u001b[0m\n\u001b[0;32m     50\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset\u001b[38;5;241m.\u001b[39m__getitems__(possibly_batched_index)\n\u001b[0;32m     51\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m---> 52\u001b[0m         data \u001b[38;5;241m=\u001b[39m [\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdataset\u001b[49m\u001b[43m[\u001b[49m\u001b[43midx\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m idx \u001b[38;5;129;01min\u001b[39;00m possibly_batched_index]\n\u001b[0;32m     53\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m     54\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[possibly_batched_index]\n",
      "File \u001b[1;32mc:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\env\\lib\\site-packages\\torchvision\\datasets\\folder.py:245\u001b[0m, in \u001b[0;36mDatasetFolder.__getitem__\u001b[1;34m(self, index)\u001b[0m\n\u001b[0;32m    237\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m    238\u001b[0m \u001b[38;5;124;03mArgs:\u001b[39;00m\n\u001b[0;32m    239\u001b[0m \u001b[38;5;124;03m    index (int): Index\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    242\u001b[0m \u001b[38;5;124;03m    tuple: (sample, target) where target is class_index of the target class.\u001b[39;00m\n\u001b[0;32m    243\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m    244\u001b[0m path, target \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msamples[index]\n\u001b[1;32m--> 245\u001b[0m sample \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mloader\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    246\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransform \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    247\u001b[0m     sample \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransform(sample)\n",
      "File \u001b[1;32mc:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\env\\lib\\site-packages\\torchvision\\datasets\\folder.py:284\u001b[0m, in \u001b[0;36mdefault_loader\u001b[1;34m(path)\u001b[0m\n\u001b[0;32m    282\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m accimage_loader(path)\n\u001b[0;32m    283\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m--> 284\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mpil_loader\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\env\\lib\\site-packages\\torchvision\\datasets\\folder.py:263\u001b[0m, in \u001b[0;36mpil_loader\u001b[1;34m(path)\u001b[0m\n\u001b[0;32m    260\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mpil_loader\u001b[39m(path: \u001b[38;5;28mstr\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Image\u001b[38;5;241m.\u001b[39mImage:\n\u001b[0;32m    261\u001b[0m     \u001b[38;5;66;03m# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)\u001b[39;00m\n\u001b[0;32m    262\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mopen\u001b[39m(path, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrb\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;28;01mas\u001b[39;00m f:\n\u001b[1;32m--> 263\u001b[0m         img \u001b[38;5;241m=\u001b[39m \u001b[43mImage\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mopen\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    264\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m img\u001b[38;5;241m.\u001b[39mconvert(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRGB\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[1;32mc:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\env\\lib\\site-packages\\PIL\\Image.py:3480\u001b[0m, in \u001b[0;36mopen\u001b[1;34m(fp, mode, formats)\u001b[0m\n\u001b[0;32m   3477\u001b[0m     fp \u001b[38;5;241m=\u001b[39m io\u001b[38;5;241m.\u001b[39mBytesIO(fp\u001b[38;5;241m.\u001b[39mread())\n\u001b[0;32m   3478\u001b[0m     exclusive_fp \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m-> 3480\u001b[0m prefix \u001b[38;5;241m=\u001b[39m \u001b[43mfp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mread\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m16\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[0;32m   3482\u001b[0m preinit()\n\u001b[0;32m   3484\u001b[0m warning_messages: \u001b[38;5;28mlist\u001b[39m[\u001b[38;5;28mstr\u001b[39m] \u001b[38;5;241m=\u001b[39m []\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "train(model, train_loader, val_loader, criterion, optimizer, num_epochs=5)\n",
    "\n",
    "# Unfreeze all the layers and fine-tune the entire network for a few more epochs\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = True\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)\n",
    "train(model, train_loader, val_loader, criterion, optimizer, num_epochs=10)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted class index: 9\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "import numpy as np\n",
    "\n",
    "# Load the trained model (assuming it's saved and restored, otherwise use the code you already have)\n",
    "model.to(\"cuda\")\n",
    "model.eval()  # Set the model to evaluation mode\n",
    "\n",
    "# Define the same transformations used during training\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "# Define a function to perform inference on an input image\n",
    "def infer(image_path, model, device):\n",
    "    # Load the image\n",
    "    img = Image.open(image_path)\n",
    "    \n",
    "    # Apply the transformations\n",
    "    img_tensor = transform(img).unsqueeze(0).to(device)  # Add batch dimension and move to the device\n",
    "    \n",
    "    # Forward pass to get predictions\n",
    "    with torch.no_grad():  # No need to track gradients during inference\n",
    "        outputs = model(img_tensor)\n",
    "        \n",
    "    # Get the predicted class\n",
    "    _, predicted_class = torch.max(outputs, 1)\n",
    "    \n",
    "    # Convert the predicted class index back to class label if necessary\n",
    "    class_idx = predicted_class.item()\n",
    "    print(f\"Predicted class index: {class_idx}\")\n",
    "    # If you have class names, you can map it to the corresponding label like this:\n",
    "    # class_names = train_dataset.classes  # The class names from your ImageFolder dataset\n",
    "    # predicted_label = class_names[class_idx]\n",
    "    # print(f\"Predicted class label: {predicted_label}\")\n",
    "\n",
    "    return class_idx  # Return the predicted class index\n",
    "\n",
    "# Example usage\n",
    "image_path = r\"C:\\Users\\Admin\\CODE\\work\\OBJECT_COLOR\\color_classify\\test_images\\0a8c3d92-3a63-4165-855d-bcb7253c67db_mask.jpg\"\n",
    "class_idx = infer(image_path, model, device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save only the model's state_dict\n",
    "torch.save(model.state_dict(), 'model_state_dict.pth')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the model again (must have the same architecture as the original)\n",
    "model = models.resnet18(pretrained=False)  # Example for ResNet-18\n",
    "\n",
    "# Modify the last layer (as in your training code)\n",
    "model.fc = torch.nn.Linear(model.fc.in_features, 13)\n",
    "\n",
    "# Load the state_dict into the model\n",
    "model.load_state_dict(torch.load('model_state_dict.pth'))\n",
    "model.eval()  # Make sure to call eval() before inference\n",
    "# Move model to device (e.g., GPU or CPU)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model.to(device)\n",
    "\n",
    "# Set the model to evaluation mode before inference\n",
    "model.eval()\n",
    "\n",
    "# Perform inference as needed\n",
    "image_path = \"path_to_your_image.jpg\"\n",
    "class_idx = infer(image_path, model, device)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}