torch_model_v1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "num_epochs = 4  \n",
    "learning_rate = 0.001\n",
    "data_dir = 'tomato'\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.Resize((128, 128)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = datasets.ImageFolder(root='tomato/train', transform=transform)\n",
    "val_dataset = datasets.ImageFolder(root='tomato/val', transform=transform)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
    "val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=32, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleCNN, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)\n",
    "        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)\n",
    "        self.fc1 = nn.Linear(128 * 16 * 16, 512)  \n",
    "        self.fc2 = nn.Linear(512, len(train_dataset.classes)) \n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = self.pool(F.relu(self.conv3(x)))\n",
    "        x = x.view(-1, 128 * 16 * 16) \n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = SimpleCNN().to(device)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10, Loss: 0.9166, Accuracy: 68.43%\n",
      "Epoch 2/10, Loss: 0.3626, Accuracy: 87.42%\n",
      "Epoch 3/10, Loss: 0.2364, Accuracy: 92.00%\n",
      "Epoch 4/10, Loss: 0.1511, Accuracy: 94.76%\n",
      "Epoch 5/10, Loss: 0.1104, Accuracy: 96.09%\n",
      "Epoch 6/10, Loss: 0.0878, Accuracy: 97.15%\n",
      "Epoch 7/10, Loss: 0.0663, Accuracy: 97.67%\n",
      "Epoch 8/10, Loss: 0.0638, Accuracy: 97.85%\n",
      "Epoch 9/10, Loss: 0.0335, Accuracy: 98.86%\n",
      "Epoch 10/10, Loss: 0.0531, Accuracy: 98.21%\n"
     ]
    }
   ],
   "source": [
    "def train_model(model, train_loader, criterion, optimizer, num_epochs=4):\n",
    "    for epoch in range(num_epochs):\n",
    "        model.train()\n",
    "        running_loss = 0.0\n",
    "        correct_predictions = 0\n",
    "        total_samples = 0\n",
    "        \n",
    "        for images, labels in train_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            running_loss += loss.item()\n",
    "            \n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            \n",
    "            total_samples += labels.size(0)\n",
    "            correct_predictions += (predicted == labels).sum().item()\n",
    "        \n",
    "        average_loss = running_loss / len(train_loader)\n",
    "        accuracy = (correct_predictions / total_samples) * 100\n",
    "        \n",
    "        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {average_loss:.4f}, Accuracy: {accuracy:.2f}%')\n",
    "        \n",
    "train_model(model, train_loader, criterion, optimizer, num_epochs=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 91.00%\n"
     ]
    }
   ],
   "source": [
    "def evaluate_model(model, val_loader):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for images, labels in val_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "    print(f'Accuracy: {100 * correct / total:.2f}%')\n",
    "\n",
    "evaluate_model(model, val_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model, './models/torch_cnn_v1.pth')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}