05_pytorch_going_modular_script_mode.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "181dd1b7-6112-4567-bef5-da71ce2ea267",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.makedirs(\"going_modular\", exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d2764fd0-ce92-4dee-ab2a-0215f6a0a566",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data/pizza_steak_sushi directory exists.\n",
      "Downloading pizza, steak, sushi data...\n",
      "Unzipping pizza, steak, sushi data...\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import zipfile\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "import requests\n",
    "\n",
    "# Setup path to data folder\n",
    "data_path = Path(\"data/\")\n",
    "image_path = data_path / \"pizza_steak_sushi\"\n",
    "\n",
    "# If the image folder doesn't exist, download it and prepare it... \n",
    "if image_path.is_dir():\n",
    "    print(f\"{image_path} directory exists.\")\n",
    "else:\n",
    "    print(f\"Did not find {image_path} directory, creating one...\")\n",
    "    image_path.mkdir(parents=True, exist_ok=True)\n",
    "    \n",
    "# Download pizza, steak, sushi data\n",
    "with open(data_path / \"pizza_steak_sushi.zip\", \"wb\") as f:\n",
    "    request = requests.get(\"https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip\")\n",
    "    print(\"Downloading pizza, steak, sushi data...\")\n",
    "    f.write(request.content)\n",
    "\n",
    "# Unzip pizza, steak, sushi data\n",
    "with zipfile.ZipFile(data_path / \"pizza_steak_sushi.zip\", \"r\") as zip_ref:\n",
    "    print(\"Unzipping pizza, steak, sushi data...\") \n",
    "    zip_ref.extractall(image_path)\n",
    "\n",
    "# Remove zip file\n",
    "os.remove(data_path / \"pizza_steak_sushi.zip\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c41b094c-ebef-4785-837c-239f3b340fad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(PosixPath('data/pizza_steak_sushi/train'),\n",
       " PosixPath('data/pizza_steak_sushi/test'))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Setup train and testing paths\n",
    "train_dir = image_path / \"train\"\n",
    "test_dir = image_path / \"test\"\n",
    "\n",
    "train_dir, test_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "574ad02b-542d-494e-9cb8-04e43e2bd5c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train data:\n",
      "Dataset ImageFolder\n",
      "    Number of datapoints: 225\n",
      "    Root location: data/pizza_steak_sushi/train\n",
      "    StandardTransform\n",
      "Transform: Compose(\n",
      "               Resize(size=(64, 64), interpolation=bilinear, max_size=None, antialias=warn)\n",
      "               ToTensor()\n",
      "           )\n",
      "Test data:\n",
      "Dataset ImageFolder\n",
      "    Number of datapoints: 75\n",
      "    Root location: data/pizza_steak_sushi/test\n",
      "    StandardTransform\n",
      "Transform: Compose(\n",
      "               Resize(size=(64, 64), interpolation=bilinear, max_size=None, antialias=warn)\n",
      "               ToTensor()\n",
      "           )\n"
     ]
    }
   ],
   "source": [
    "from torchvision import datasets, transforms\n",
    "\n",
    "# Create simple transform\n",
    "data_transform = transforms.Compose([ \n",
    "    transforms.Resize((64, 64)),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "# Use ImageFolder to create dataset(s)\n",
    "train_data = datasets.ImageFolder(root=train_dir, # target folder of images\n",
    "                                  transform=data_transform, # transforms to perform on data (images)\n",
    "                                  target_transform=None) # transforms to perform on labels (if necessary)\n",
    "\n",
    "test_data = datasets.ImageFolder(root=test_dir, \n",
    "                                 transform=data_transform)\n",
    "\n",
    "print(f\"Train data:\\n{train_data}\\nTest data:\\n{test_data}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d02ff25f-1227-4e6e-9615-2d350089b1fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['pizza', 'steak', 'sushi']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get class names as a list\n",
    "class_names = train_data.classes\n",
    "class_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1a8dc077-8b55-47a6-81bd-899985d134bb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'pizza': 0, 'steak': 1, 'sushi': 2}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Can also get class names as a dict\n",
    "class_dict = train_data.class_to_idx\n",
    "class_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6ae929ae-6945-4b0e-9ff8-979387a366d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(225, 75)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check the lengths\n",
    "len(train_data), len(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4cd8f99a-4f58-497d-b4b0-fe3df4311597",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<torch.utils.data.dataloader.DataLoader at 0x73198e95c4d0>,\n",
       " <torch.utils.data.dataloader.DataLoader at 0x7319900b5350>)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Turn train and test Datasets into DataLoaders\n",
    "from torch.utils.data import DataLoader\n",
    "NUMBER_WORKERS = os.cpu_count()\n",
    "train_dataloader = DataLoader(dataset=train_data, \n",
    "                              batch_size=32, # how many samples per batch?\n",
    "                              num_workers=NUMBER_WORKERS, # how many subprocesses to use for data loading? (higher = more)\n",
    "                              shuffle=True) # shuffle the data?\n",
    "\n",
    "test_dataloader = DataLoader(dataset=test_data, \n",
    "                             batch_size=32, \n",
    "                             num_workers=NUMBER_WORKERS, \n",
    "                             shuffle=False) # don't usually need to shuffle testing data\n",
    "\n",
    "train_dataloader, test_dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "79251739-0f7c-41f9-b530-b68c889ebfae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image shape: torch.Size([32, 3, 64, 64]) -> [batch_size, color_channels, height, width]\n",
      "Label shape: torch.Size([32])\n"
     ]
    }
   ],
   "source": [
    "# Check out single image size/shape\n",
    "img, label = next(iter(train_dataloader))\n",
    "\n",
    "# Batch size will now be 1, try changing the batch_size parameter above and see what happens\n",
    "print(f\"Image shape: {img.shape} -> [batch_size, color_channels, height, width]\")\n",
    "print(f\"Label shape: {label.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a0f3be30-9d8a-4a17-b99b-f64142f6ac4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting going_modular/data_setup.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile going_modular/data_setup.py\n",
    "\"\"\"\n",
    "Contains functionality for creating PyTorch DataLoaders for \n",
    "image classification data.\n",
    "\"\"\"\n",
    "import os\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "NUM_WORKERS = os.cpu_count()\n",
    "\n",
    "def create_dataloaders(\n",
    "    train_dir: str, \n",
    "    test_dir: str, \n",
    "    transform: transforms.Compose, \n",
    "    batch_size: int, \n",
    "    num_workers: int=NUM_WORKERS\n",
    "):\n",
    "  \"\"\"Creates training and testing DataLoaders.\n",
    "\n",
    "  Takes in a training directory and testing directory path and turns\n",
    "  them into PyTorch Datasets and then into PyTorch DataLoaders.\n",
    "\n",
    "  Args:\n",
    "    train_dir: Path to training directory.\n",
    "    test_dir: Path to testing directory.\n",
    "    transform: torchvision transforms to perform on training and testing data.\n",
    "    batch_size: Number of samples per batch in each of the DataLoaders.\n",
    "    num_workers: An integer for number of workers per DataLoader.\n",
    "\n",
    "  Returns:\n",
    "    A tuple of (train_dataloader, test_dataloader, class_names).\n",
    "    Where class_names is a list of the target classes.\n",
    "    Example usage:\n",
    "      train_dataloader, test_dataloader, class_names = \\\n",
    "        = create_dataloaders(train_dir=path/to/train_dir,\n",
    "                             test_dir=path/to/test_dir,\n",
    "                             transform=some_transform,\n",
    "                             batch_size=32,\n",
    "                             num_workers=4)\n",
    "  \"\"\"\n",
    "  # Use ImageFolder to create dataset(s)\n",
    "  train_data = datasets.ImageFolder(train_dir, transform=transform)\n",
    "  test_data = datasets.ImageFolder(test_dir, transform=transform)\n",
    "\n",
    "  # Get class names\n",
    "  class_names = train_data.classes\n",
    "\n",
    "  # Turn images into data loaders\n",
    "  train_dataloader = DataLoader(\n",
    "      train_data,\n",
    "      batch_size=batch_size,\n",
    "      shuffle=True,\n",
    "      num_workers=num_workers,\n",
    "      pin_memory=True,\n",
    "  )\n",
    "  test_dataloader = DataLoader(\n",
    "      test_data,\n",
    "      batch_size=batch_size,\n",
    "      shuffle=False,\n",
    "      num_workers=num_workers,\n",
    "      pin_memory=True,\n",
    "  )\n",
    "\n",
    "  return train_dataloader, test_dataloader, class_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6e6a7b52-cea0-45b4-9b9d-f4085ccb8c3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import data_setup.py\n",
    "from going_modular import data_setup\n",
    "\n",
    "# Create train/test dataloader and get class names as a list\n",
    "train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(train_dir, test_dir, data_transform, 32, NUMBER_WORKERS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "980da930-953c-44c2-b4b6-d58ca196965b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['pizza', 'steak', 'sushi']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "955809ba-e7b0-478b-895c-7db2460790a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from torch import nn \n",
    "\n",
    "class TinyVGG(nn.Module):\n",
    "    \"\"\"Creates the TinyVGG architecture.\n",
    "\n",
    "    Replicates the TinyVGG architecture from the CNN explainer website in PyTorch.\n",
    "    See the original architecture here: https://poloclub.github.io/cnn-explainer/\n",
    "\n",
    "    Args:\n",
    "    input_shape: An integer indicating number of input channels.\n",
    "    hidden_units: An integer indicating number of hidden units between layers.\n",
    "    output_shape: An integer indicating number of output units.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:\n",
    "        super().__init__()\n",
    "        self.conv_block_1 = nn.Sequential(\n",
    "          nn.Conv2d(in_channels=input_shape, \n",
    "                    out_channels=hidden_units, \n",
    "                    kernel_size=3, \n",
    "                    stride=1, \n",
    "                    padding=0),  \n",
    "          nn.ReLU(),\n",
    "          nn.Conv2d(in_channels=hidden_units, \n",
    "                    out_channels=hidden_units,\n",
    "                    kernel_size=3,\n",
    "                    stride=1,\n",
    "                    padding=0),\n",
    "          nn.ReLU(),\n",
    "          nn.MaxPool2d(kernel_size=2,\n",
    "                        stride=2)\n",
    "        )\n",
    "        self.conv_block_2 = nn.Sequential(\n",
    "          nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),\n",
    "          nn.ReLU(),\n",
    "          nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),\n",
    "          nn.ReLU(),\n",
    "          nn.MaxPool2d(2)\n",
    "        )\n",
    "        self.classifier = nn.Sequential(\n",
    "          nn.Flatten(),\n",
    "          # Where did this in_features shape come from? \n",
    "          # It's because each layer of our network compresses and changes the shape of our inputs data.\n",
    "          nn.Linear(in_features=hidden_units*13*13,\n",
    "                    out_features=output_shape)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x: torch.Tensor):\n",
    "        x = self.conv_block_1(x)\n",
    "        x = self.conv_block_2(x)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "        # return self.classifier(self.block_2(self.block_1(x))) # <- leverage the benefits of operator fusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "87b0cf63-7f4a-4271-a6f6-2e6e0a1632e2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TinyVGG(\n",
       "  (conv_block_1): Sequential(\n",
       "    (0): Conv2d(3, 10, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): ReLU()\n",
       "    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (3): ReLU()\n",
       "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  )\n",
       "  (conv_block_2): Sequential(\n",
       "    (0): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): ReLU()\n",
       "    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (3): ReLU()\n",
       "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  )\n",
       "  (classifier): Sequential(\n",
       "    (0): Flatten(start_dim=1, end_dim=-1)\n",
       "    (1): Linear(in_features=1690, out_features=3, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# Instantiate an instance of the model\n",
    "torch.manual_seed(42)\n",
    "model_0 = TinyVGG(input_shape=3, # number of color channels (3 for RGB) \n",
    "                  hidden_units=10, \n",
    "                  output_shape=len(train_data.classes)).to(device)\n",
    "model_0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "49828e94-51fe-45e4-b476-c33915b509e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Single image shape: torch.Size([1, 3, 64, 64])\n",
      "\n",
      "Output logits:\n",
      "tensor([[ 0.0208, -0.0020,  0.0095]], device='cuda:0')\n",
      "\n",
      "Output prediction probabilities:\n",
      "tensor([[0.3371, 0.3295, 0.3333]], device='cuda:0')\n",
      "\n",
      "Output prediction label:\n",
      "tensor([0], device='cuda:0')\n",
      "\n",
      "Actual label:\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "# 1. Get a batch of images and labels from the DataLoader\n",
    "img_batch, label_batch = next(iter(train_dataloader))\n",
    "\n",
    "# 2. Get a single image from the batch and unsqueeze the image so its shape fits the model\n",
    "img_single, label_single = img_batch[0].unsqueeze(dim=0), label_batch[0]\n",
    "print(f\"Single image shape: {img_single.shape}\\n\")\n",
    "\n",
    "# 3. Perform a forward pass on a single image\n",
    "model_0.eval()\n",
    "with torch.inference_mode():\n",
    "    pred = model_0(img_single.to(device))\n",
    "    \n",
    "# 4. Print out what's happening and convert model logits -> pred probs -> pred label\n",
    "print(f\"Output logits:\\n{pred}\\n\")\n",
    "print(f\"Output prediction probabilities:\\n{torch.softmax(pred, dim=1)}\\n\")\n",
    "print(f\"Output prediction label:\\n{torch.argmax(torch.softmax(pred, dim=1), dim=1)}\\n\")\n",
    "print(f\"Actual label:\\n{label_single}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d07699dc-88db-4eb4-917c-c542d8ba5766",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting going_modular/model_builder.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile going_modular/model_builder.py\n",
    "\"\"\"\n",
    "Contains PyTorch model code to instantiate a TinyVGG model.\n",
    "\"\"\"\n",
    "import torch\n",
    "\n",
    "from torch import nn\n",
    "\n",
    "class TinyVGG(nn.Module):\n",
    "    \"\"\"Creates the TinyVGG architecture.\n",
    "\n",
    "    Replicates the TinyVGG architecture from the CNN explainer website in PyTorch.\n",
    "    See the original architecture here: https://poloclub.github.io/cnn-explainer/\n",
    "\n",
    "    Args:\n",
    "    input_shape: An integer indicating number of input channels.\n",
    "    hidden_units: An integer indicating number of hidden units between layers.\n",
    "    output_shape: An integer indicating number of output units.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:\n",
    "        super().__init__()\n",
    "        self.conv_block_1 = nn.Sequential(\n",
    "          nn.Conv2d(in_channels=input_shape, \n",
    "                    out_channels=hidden_units, \n",
    "                    kernel_size=3, \n",
    "                    stride=1, \n",
    "                    padding=0),  \n",
    "          nn.ReLU(),\n",
    "          nn.Conv2d(in_channels=hidden_units, \n",
    "                    out_channels=hidden_units,\n",
    "                    kernel_size=3,\n",
    "                    stride=1,\n",
    "                    padding=0),\n",
    "          nn.ReLU(),\n",
    "          nn.MaxPool2d(kernel_size=2,\n",
    "                        stride=2)\n",
    "        )\n",
    "        self.conv_block_2 = nn.Sequential(\n",
    "          nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),\n",
    "          nn.ReLU(),\n",
    "          nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),\n",
    "          nn.ReLU(),\n",
    "          nn.MaxPool2d(2)\n",
    "        )\n",
    "        self.classifier = nn.Sequential(\n",
    "          nn.Flatten(),\n",
    "          # Where did this in_features shape come from? \n",
    "          # It's because each layer of our network compresses and changes the shape of our inputs data.\n",
    "          nn.Linear(in_features=hidden_units*13*13,\n",
    "                    out_features=output_shape)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x: torch.Tensor):\n",
    "        #x = self.conv_block_1(x)\n",
    "        #x = self.conv_block_2(x)\n",
    "        #x = self.classifier(x)\n",
    "        #return x\n",
    "        return self.classifier(self.conv_block_2(self.conv_block_1(x))) # <- leverage the benefits of operator fusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "2d4e00e9-482f-4ba4-ad16-52aac6e9ca82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TinyVGG(\n",
       "  (conv_block_1): Sequential(\n",
       "    (0): Conv2d(3, 10, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): ReLU()\n",
       "    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (3): ReLU()\n",
       "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  )\n",
       "  (conv_block_2): Sequential(\n",
       "    (0): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): ReLU()\n",
       "    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (3): ReLU()\n",
       "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  )\n",
       "  (classifier): Sequential(\n",
       "    (0): Flatten(start_dim=1, end_dim=-1)\n",
       "    (1): Linear(in_features=1690, out_features=3, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "from going_modular import model_builder\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# Instantiate an instance of the model from the \"model_builder.py\" script\n",
    "torch.manual_seed(42)\n",
    "model_1 = model_builder.TinyVGG(input_shape=3, # number of color channels (3 for RGB) \n",
    "                                hidden_units=10, \n",
    "                                output_shape=len(class_names)).to(device)\n",
    "model_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1453ae9e-3556-41f1-9ed7-efc0b9e80101",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Single image shape: torch.Size([1, 3, 64, 64])\n",
      "\n",
      "Output logits:\n",
      "tensor([[ 0.0208, -0.0020,  0.0095]], device='cuda:0')\n",
      "\n",
      "Output prediction probabilities:\n",
      "tensor([[0.3371, 0.3295, 0.3333]], device='cuda:0')\n",
      "\n",
      "Output prediction label:\n",
      "tensor([0], device='cuda:0')\n",
      "\n",
      "Actual label:\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "# 1. Get a batch of images and labels from the DataLoader\n",
    "img_batch, label_batch = next(iter(train_dataloader))\n",
    "\n",
    "# 2. Get a single image from the batch and unsqueeze the image so its shape fits the model\n",
    "img_single, label_single = img_batch[0].unsqueeze(dim=0), label_batch[0]\n",
    "print(f\"Single image shape: {img_single.shape}\\n\")\n",
    "\n",
    "# 3. Perform a forward pass on a single image\n",
    "model_1.eval()\n",
    "with torch.inference_mode():\n",
    "    pred = model_1(img_single.to(device))\n",
    "    \n",
    "# 4. Print out what's happening and convert model logits -> pred probs -> pred label\n",
    "print(f\"Output logits:\\n{pred}\\n\")\n",
    "print(f\"Output prediction probabilities:\\n{torch.softmax(pred, dim=1)}\\n\")\n",
    "print(f\"Output prediction label:\\n{torch.argmax(torch.softmax(pred, dim=1), dim=1)}\\n\")\n",
    "print(f\"Actual label:\\n{label_single}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "988255b3-a3ea-48c3-905a-e5135b95d248",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple\n",
    "\n",
    "def train_step(model: torch.nn.Module, \n",
    "               dataloader: torch.utils.data.DataLoader, \n",
    "               loss_fn: torch.nn.Module, \n",
    "               optimizer: torch.optim.Optimizer,\n",
    "               device: torch.device) -> Tuple[float, float]:\n",
    "    \"\"\"Trains a PyTorch model for a single epoch.\n",
    "\n",
    "    Turns a target PyTorch model to training mode and then\n",
    "    runs through all of the required training steps (forward\n",
    "    pass, loss calculation, optimizer step).\n",
    "\n",
    "    Args:\n",
    "    model: A PyTorch model to be trained.\n",
    "    dataloader: A DataLoader instance for the model to be trained on.\n",
    "    loss_fn: A PyTorch loss function to minimize.\n",
    "    optimizer: A PyTorch optimizer to help minimize the loss function.\n",
    "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "    A tuple of training loss and training accuracy metrics.\n",
    "    In the form (train_loss, train_accuracy). For example:\n",
    "\n",
    "    (0.1112, 0.8743)\n",
    "    \"\"\"\n",
    "    # Put model in train mode\n",
    "    model.train()\n",
    "\n",
    "    # Setup train loss and train accuracy values\n",
    "    train_loss, train_acc = 0, 0\n",
    "\n",
    "    # Loop through data loader data batches\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        # Send data to target device\n",
    "        X, y = X.to(device), y.to(device)\n",
    "\n",
    "        # 1. Forward pass\n",
    "        y_pred = model(X)\n",
    "\n",
    "        # 2. Calculate  and accumulate loss\n",
    "        loss = loss_fn(y_pred, y)\n",
    "        train_loss += loss.item() \n",
    "\n",
    "        # 3. Optimizer zero grad\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # 4. Loss backward\n",
    "        loss.backward()\n",
    "\n",
    "        # 5. Optimizer step\n",
    "        optimizer.step()\n",
    "\n",
    "        # Calculate and accumulate accuracy metric across all batches\n",
    "        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)\n",
    "        train_acc += (y_pred_class == y).sum().item()/len(y_pred)\n",
    "\n",
    "    # Adjust metrics to get average loss and accuracy per batch \n",
    "    train_loss = train_loss / len(dataloader)\n",
    "    train_acc = train_acc / len(dataloader)\n",
    "    return train_loss, train_acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3e030cb4-1fc7-4693-abc1-2a3e3be50f4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_step(model: torch.nn.Module, \n",
    "              dataloader: torch.utils.data.DataLoader, \n",
    "              loss_fn: torch.nn.Module,\n",
    "              device: torch.device) -> Tuple[float, float]:\n",
    "    \"\"\"Tests a PyTorch model for a single epoch.\n",
    "\n",
    "    Turns a target PyTorch model to \"eval\" mode and then performs\n",
    "    a forward pass on a testing dataset.\n",
    "\n",
    "    Args:\n",
    "    model: A PyTorch model to be tested.\n",
    "    dataloader: A DataLoader instance for the model to be tested on.\n",
    "    loss_fn: A PyTorch loss function to calculate loss on the test data.\n",
    "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "    A tuple of testing loss and testing accuracy metrics.\n",
    "    In the form (test_loss, test_accuracy). For example:\n",
    "\n",
    "    (0.0223, 0.8985)\n",
    "    \"\"\"\n",
    "    # Put model in eval mode\n",
    "    model.eval() \n",
    "\n",
    "    # Setup test loss and test accuracy values\n",
    "    test_loss, test_acc = 0, 0\n",
    "\n",
    "    # Turn on inference context manager\n",
    "    with torch.inference_mode():\n",
    "        # Loop through DataLoader batches\n",
    "        for batch, (X, y) in enumerate(dataloader):\n",
    "            # Send data to target device\n",
    "            X, y = X.to(device), y.to(device)\n",
    "\n",
    "            # 1. Forward pass\n",
    "            test_pred_logits = model(X)\n",
    "\n",
    "            # 2. Calculate and accumulate loss\n",
    "            loss = loss_fn(test_pred_logits, y)\n",
    "            test_loss += loss.item()\n",
    "\n",
    "            # Calculate and accumulate accuracy\n",
    "            test_pred_labels = test_pred_logits.argmax(dim=1)\n",
    "            test_acc += ((test_pred_labels == y).sum().item()/len(test_pred_labels))\n",
    "\n",
    "    # Adjust metrics to get average loss and accuracy per batch \n",
    "    test_loss = test_loss / len(dataloader)\n",
    "    test_acc = test_acc / len(dataloader)\n",
    "    return test_loss, test_acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e7e60743-d413-49a8-8b88-4d4d647af1e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict, List\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "def train(model: torch.nn.Module, \n",
    "          train_dataloader: torch.utils.data.DataLoader, \n",
    "          test_dataloader: torch.utils.data.DataLoader, \n",
    "          optimizer: torch.optim.Optimizer,\n",
    "          loss_fn: torch.nn.Module,\n",
    "          epochs: int,\n",
    "          device: torch.device) -> Dict[str, List[float]]:\n",
    "    \"\"\"Trains and tests a PyTorch model.\n",
    "\n",
    "    Passes a target PyTorch models through train_step() and test_step()\n",
    "    functions for a number of epochs, training and testing the model\n",
    "    in the same epoch loop.\n",
    "\n",
    "    Calculates, prints and stores evaluation metrics throughout.\n",
    "\n",
    "    Args:\n",
    "    model: A PyTorch model to be trained and tested.\n",
    "    train_dataloader: A DataLoader instance for the model to be trained on.\n",
    "    test_dataloader: A DataLoader instance for the model to be tested on.\n",
    "    optimizer: A PyTorch optimizer to help minimize the loss function.\n",
    "    loss_fn: A PyTorch loss function to calculate loss on both datasets.\n",
    "    epochs: An integer indicating how many epochs to train for.\n",
    "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "    A dictionary of training and testing loss as well as training and\n",
    "    testing accuracy metrics. Each metric has a value in a list for \n",
    "    each epoch.\n",
    "    In the form: {train_loss: [...],\n",
    "                  train_acc: [...],\n",
    "                  test_loss: [...],\n",
    "                  test_acc: [...]} \n",
    "    For example if training for epochs=2: \n",
    "                 {train_loss: [2.0616, 1.0537],\n",
    "                  train_acc: [0.3945, 0.3945],\n",
    "                  test_loss: [1.2641, 1.5706],\n",
    "                  test_acc: [0.3400, 0.2973]} \n",
    "    \"\"\"\n",
    "    # Create empty results dictionary\n",
    "    results = {\"train_loss\": [],\n",
    "      \"train_acc\": [],\n",
    "      \"test_loss\": [],\n",
    "      \"test_acc\": []\n",
    "    }\n",
    "\n",
    "    # Loop through training and testing steps for a number of epochs\n",
    "    for epoch in tqdm(range(epochs)):\n",
    "        train_loss, train_acc = train_step(model=model,\n",
    "                                          dataloader=train_dataloader,\n",
    "                                          loss_fn=loss_fn,\n",
    "                                          optimizer=optimizer,\n",
    "                                          device=device)\n",
    "        test_loss, test_acc = test_step(model=model,\n",
    "          dataloader=test_dataloader,\n",
    "          loss_fn=loss_fn,\n",
    "          device=device)\n",
    "\n",
    "        # Print out what's happening\n",
    "        print(\n",
    "          f\"Epoch: {epoch+1} | \"\n",
    "          f\"train_loss: {train_loss:.4f} | \"\n",
    "          f\"train_acc: {train_acc:.4f} | \"\n",
    "          f\"test_loss: {test_loss:.4f} | \"\n",
    "          f\"test_acc: {test_acc:.4f}\"\n",
    "        )\n",
    "\n",
    "        # Update results dictionary\n",
    "        results[\"train_loss\"].append(train_loss)\n",
    "        results[\"train_acc\"].append(train_acc)\n",
    "        results[\"test_loss\"].append(test_loss)\n",
    "        results[\"test_acc\"].append(test_acc)\n",
    "\n",
    "    # Return the filled results at the end of the epochs\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "7e25cc1c-3070-43e7-a5de-4f5b0f4f04e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing going_modular/engine.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile going_modular/engine.py\n",
    "\"\"\"\n",
    "Contains functions for training and testing a PyTorch model.\n",
    "\"\"\"\n",
    "from typing import Dict, List, Tuple\n",
    "\n",
    "import torch\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "def train_step(model: torch.nn.Module, \n",
    "               dataloader: torch.utils.data.DataLoader, \n",
    "               loss_fn: torch.nn.Module, \n",
    "               optimizer: torch.optim.Optimizer,\n",
    "               device: torch.device) -> Tuple[float, float]:\n",
    "    \"\"\"Trains a PyTorch model for a single epoch.\n",
    "\n",
    "    Turns a target PyTorch model to training mode and then\n",
    "    runs through all of the required training steps (forward\n",
    "    pass, loss calculation, optimizer step).\n",
    "\n",
    "    Args:\n",
    "    model: A PyTorch model to be trained.\n",
    "    dataloader: A DataLoader instance for the model to be trained on.\n",
    "    loss_fn: A PyTorch loss function to minimize.\n",
    "    optimizer: A PyTorch optimizer to help minimize the loss function.\n",
    "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "    A tuple of training loss and training accuracy metrics.\n",
    "    In the form (train_loss, train_accuracy). For example:\n",
    "\n",
    "    (0.1112, 0.8743)\n",
    "    \"\"\"\n",
    "    # Put model in train mode\n",
    "    model.train()\n",
    "\n",
    "    # Setup train loss and train accuracy values\n",
    "    train_loss, train_acc = 0, 0\n",
    "\n",
    "    # Loop through data loader data batches\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        # Send data to target device\n",
    "        X, y = X.to(device), y.to(device)\n",
    "\n",
    "        # 1. Forward pass\n",
    "        y_pred = model(X)\n",
    "\n",
    "        # 2. Calculate  and accumulate loss\n",
    "        loss = loss_fn(y_pred, y)\n",
    "        train_loss += loss.item() \n",
    "\n",
    "        # 3. Optimizer zero grad\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # 4. Loss backward\n",
    "        loss.backward()\n",
    "\n",
    "        # 5. Optimizer step\n",
    "        optimizer.step()\n",
    "\n",
    "        # Calculate and accumulate accuracy metric across all batches\n",
    "        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)\n",
    "        train_acc += (y_pred_class == y).sum().item()/len(y_pred)\n",
    "\n",
    "    # Adjust metrics to get average loss and accuracy per batch \n",
    "    train_loss = train_loss / len(dataloader)\n",
    "    train_acc = train_acc / len(dataloader)\n",
    "    return train_loss, train_acc\n",
    "\n",
    "def test_step(model: torch.nn.Module, \n",
    "              dataloader: torch.utils.data.DataLoader, \n",
    "              loss_fn: torch.nn.Module,\n",
    "              device: torch.device) -> Tuple[float, float]:\n",
    "    \"\"\"Tests a PyTorch model for a single epoch.\n",
    "\n",
    "    Turns a target PyTorch model to \"eval\" mode and then performs\n",
    "    a forward pass on a testing dataset.\n",
    "\n",
    "    Args:\n",
    "    model: A PyTorch model to be tested.\n",
    "    dataloader: A DataLoader instance for the model to be tested on.\n",
    "    loss_fn: A PyTorch loss function to calculate loss on the test data.\n",
    "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "    A tuple of testing loss and testing accuracy metrics.\n",
    "    In the form (test_loss, test_accuracy). For example:\n",
    "\n",
    "    (0.0223, 0.8985)\n",
    "    \"\"\"\n",
    "    # Put model in eval mode\n",
    "    model.eval() \n",
    "\n",
    "    # Setup test loss and test accuracy values\n",
    "    test_loss, test_acc = 0, 0\n",
    "\n",
    "    # Turn on inference context manager\n",
    "    with torch.inference_mode():\n",
    "        # Loop through DataLoader batches\n",
    "        for batch, (X, y) in enumerate(dataloader):\n",
    "            # Send data to target device\n",
    "            X, y = X.to(device), y.to(device)\n",
    "\n",
    "            # 1. Forward pass\n",
    "            test_pred_logits = model(X)\n",
    "\n",
    "            # 2. Calculate and accumulate loss\n",
    "            loss = loss_fn(test_pred_logits, y)\n",
    "            test_loss += loss.item()\n",
    "\n",
    "            # Calculate and accumulate accuracy\n",
    "            test_pred_labels = test_pred_logits.argmax(dim=1)\n",
    "            test_acc += ((test_pred_labels == y).sum().item()/len(test_pred_labels))\n",
    "\n",
    "    # Adjust metrics to get average loss and accuracy per batch \n",
    "    test_loss = test_loss / len(dataloader)\n",
    "    test_acc = test_acc / len(dataloader)\n",
    "    return test_loss, test_acc\n",
    "\n",
    "def train(model: torch.nn.Module, \n",
    "          train_dataloader: torch.utils.data.DataLoader, \n",
    "          test_dataloader: torch.utils.data.DataLoader, \n",
    "          optimizer: torch.optim.Optimizer,\n",
    "          loss_fn: torch.nn.Module,\n",
    "          epochs: int,\n",
    "          device: torch.device) -> Dict[str, List[float]]:\n",
    "    \"\"\"Trains and tests a PyTorch model.\n",
    "\n",
    "    Passes a target PyTorch models through train_step() and test_step()\n",
    "    functions for a number of epochs, training and testing the model\n",
    "    in the same epoch loop.\n",
    "\n",
    "    Calculates, prints and stores evaluation metrics throughout.\n",
    "\n",
    "    Args:\n",
    "    model: A PyTorch model to be trained and tested.\n",
    "    train_dataloader: A DataLoader instance for the model to be trained on.\n",
    "    test_dataloader: A DataLoader instance for the model to be tested on.\n",
    "    optimizer: A PyTorch optimizer to help minimize the loss function.\n",
    "    loss_fn: A PyTorch loss function to calculate loss on both datasets.\n",
    "    epochs: An integer indicating how many epochs to train for.\n",
    "    device: A target device to compute on (e.g. \"cuda\" or \"cpu\").\n",
    "\n",
    "    Returns:\n",
    "    A dictionary of training and testing loss as well as training and\n",
    "    testing accuracy metrics. Each metric has a value in a list for \n",
    "    each epoch.\n",
    "    In the form: {train_loss: [...],\n",
    "              train_acc: [...],\n",
    "              test_loss: [...],\n",
    "              test_acc: [...]} \n",
    "    For example if training for epochs=2: \n",
    "             {train_loss: [2.0616, 1.0537],\n",
    "              train_acc: [0.3945, 0.3945],\n",
    "              test_loss: [1.2641, 1.5706],\n",
    "              test_acc: [0.3400, 0.2973]} \n",
    "    \"\"\"\n",
    "    # Create empty results dictionary\n",
    "    results = {\"train_loss\": [],\n",
    "               \"train_acc\": [],\n",
    "               \"test_loss\": [],\n",
    "               \"test_acc\": []\n",
    "    }\n",
    "\n",
    "    # Loop through training and testing steps for a number of epochs\n",
    "    for epoch in tqdm(range(epochs)):\n",
    "        train_loss, train_acc = train_step(model=model,\n",
    "                                          dataloader=train_dataloader,\n",
    "                                          loss_fn=loss_fn,\n",
    "                                          optimizer=optimizer,\n",
    "                                          device=device)\n",
    "        test_loss, test_acc = test_step(model=model,\n",
    "          dataloader=test_dataloader,\n",
    "          loss_fn=loss_fn,\n",
    "          device=device)\n",
    "\n",
    "        # Print out what's happening\n",
    "        print(\n",
    "          f\"Epoch: {epoch+1} | \"\n",
    "          f\"train_loss: {train_loss:.4f} | \"\n",
    "          f\"train_acc: {train_acc:.4f} | \"\n",
    "          f\"test_loss: {test_loss:.4f} | \"\n",
    "          f\"test_acc: {test_acc:.4f}\"\n",
    "        )\n",
    "\n",
    "        # Update results dictionary\n",
    "        results[\"train_loss\"].append(train_loss)\n",
    "        results[\"train_acc\"].append(train_acc)\n",
    "        results[\"test_loss\"].append(test_loss)\n",
    "        results[\"test_acc\"].append(test_acc)\n",
    "\n",
    "    # Return the filled results at the end of the epochs\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b842da55-c8af-4bd0-8688-cbae44de84cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "def save_model(model: torch.nn.Module,\n",
    "               target_dir: str,\n",
    "               model_name: str):\n",
    "    \"\"\"Saves a PyTorch model to a target directory.\n",
    "\n",
    "    Args:\n",
    "    model: A target PyTorch model to save.\n",
    "    target_dir: A directory for saving the model to.\n",
    "    model_name: A filename for the saved model. Should include\n",
    "      either \".pth\" or \".pt\" as the file extension.\n",
    "\n",
    "    Example usage:\n",
    "    save_model(model=model_0,\n",
    "               target_dir=\"models\",\n",
    "               model_name=\"05_going_modular_tingvgg_model.pth\")\n",
    "    \"\"\"\n",
    "    # Create target directory\n",
    "    target_dir_path = Path(target_dir)\n",
    "    target_dir_path.mkdir(parents=True,\n",
    "                        exist_ok=True)\n",
    "\n",
    "    # Create model save path\n",
    "    assert model_name.endswith(\".pth\") or model_name.endswith(\".pt\"), \"model_name should end with '.pt' or '.pth'\"\n",
    "    model_save_path = target_dir_path / model_name\n",
    "\n",
    "    # Save the model state_dict()\n",
    "    print(f\"[INFO] Saving model to: {model_save_path}\")\n",
    "    torch.save(obj=model.state_dict(),\n",
    "             f=model_save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "9729849d-8c9b-47fc-8002-10487900222f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing going_modular/utils.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile going_modular/utils.py\n",
    "\"\"\"\n",
    "Contains various utility functions for PyTorch model training and saving.\n",
    "\"\"\"\n",
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "\n",
    "def save_model(model: torch.nn.Module,\n",
    "               target_dir: str,\n",
    "               model_name: str):\n",
    "    \"\"\"Saves a PyTorch model to a target directory.\n",
    "\n",
    "    Args:\n",
    "    model: A target PyTorch model to save.\n",
    "    target_dir: A directory for saving the model to.\n",
    "    model_name: A filename for the saved model. Should include\n",
    "      either \".pth\" or \".pt\" as the file extension.\n",
    "\n",
    "    Example usage:\n",
    "    save_model(model=model_0,\n",
    "               target_dir=\"models\",\n",
    "               model_name=\"05_going_modular_tingvgg_model.pth\")\n",
    "    \"\"\"\n",
    "    # Create target directory\n",
    "    target_dir_path = Path(target_dir)\n",
    "    target_dir_path.mkdir(parents=True,\n",
    "                        exist_ok=True)\n",
    "\n",
    "    # Create model save path\n",
    "    assert model_name.endswith(\".pth\") or model_name.endswith(\".pt\"), \"model_name should end with '.pt' or '.pth'\"\n",
    "    model_save_path = target_dir_path / model_name\n",
    "\n",
    "    # Save the model state_dict()\n",
    "    print(f\"[INFO] Saving model to: {model_save_path}\")\n",
    "    torch.save(obj=model.state_dict(),\n",
    "             f=model_save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8b888929-c512-48d0-bfff-1f0530465007",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fee7c30180ec49bca931f2f27b1a52cb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1 | train_loss: 1.1063 | train_acc: 0.3047 | test_loss: 1.0983 | test_acc: 0.3011\n",
      "Epoch: 2 | train_loss: 1.0998 | train_acc: 0.3281 | test_loss: 1.0697 | test_acc: 0.5417\n",
      "Epoch: 3 | train_loss: 1.0869 | train_acc: 0.4883 | test_loss: 1.0807 | test_acc: 0.4924\n",
      "Epoch: 4 | train_loss: 1.0847 | train_acc: 0.4023 | test_loss: 1.0605 | test_acc: 0.5833\n",
      "Epoch: 5 | train_loss: 1.0663 | train_acc: 0.4180 | test_loss: 1.0657 | test_acc: 0.5644\n",
      "[INFO] Total training time: 2.493 seconds\n",
      "[INFO] Saving model to: models/05_going_modular_cell_mode_tinyvgg_model.pth\n"
     ]
    }
   ],
   "source": [
    "# Set random seeds\n",
    "torch.manual_seed(42) \n",
    "torch.cuda.manual_seed(42)\n",
    "\n",
    "# Set number of epochs\n",
    "NUM_EPOCHS = 5\n",
    "\n",
    "# Recreate an instance of TinyVGG\n",
    "model_0 = TinyVGG(input_shape=3, # number of color channels (3 for RGB) \n",
    "                  hidden_units=10, \n",
    "                  output_shape=len(train_data.classes)).to(device)\n",
    "\n",
    "# Setup loss function and optimizer\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(params=model_0.parameters(), lr=0.001)\n",
    "\n",
    "# Start the timer\n",
    "from timeit import default_timer as timer \n",
    "start_time = timer()\n",
    "\n",
    "# Train model_0 \n",
    "model_0_results = train(model=model_0, \n",
    "                        train_dataloader=train_dataloader,\n",
    "                        test_dataloader=test_dataloader,\n",
    "                        optimizer=optimizer,\n",
    "                        loss_fn=loss_fn, \n",
    "                        epochs=NUM_EPOCHS,\n",
    "                        device=device)\n",
    "\n",
    "# End the timer and print out how long it took\n",
    "end_time = timer()\n",
    "print(f\"[INFO] Total training time: {end_time-start_time:.3f} seconds\")\n",
    "\n",
    "# Save the model\n",
    "save_model(model=model_0,\n",
    "           target_dir=\"models\",\n",
    "           model_name=\"05_going_modular_cell_mode_tinyvgg_model.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "30b08875-89fb-4cc2-851a-927b24b9a4b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing going_modular/train.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile going_modular/train.py\n",
    "\"\"\"\n",
    "Trains a PyTorch image classification model using device-agnostic code.\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "\n",
    "import torch\n",
    "\n",
    "from torchvision import transforms\n",
    "\n",
    "import data_setup, engine, model_builder, utils\n",
    "\n",
    "# Setup hyperparameters\n",
    "NUM_EPOCHS = 5\n",
    "BATCH_SIZE = 32\n",
    "HIDDEN_UNITS = 10\n",
    "LEARNING_RATE = 0.001\n",
    "\n",
    "# Setup directories\n",
    "train_dir = \"data/pizza_steak_sushi/train\"\n",
    "test_dir = \"data/pizza_steak_sushi/test\"\n",
    "\n",
    "# Setup target device\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# Create transforms\n",
    "data_transform = transforms.Compose([\n",
    "  transforms.Resize((64, 64)),\n",
    "  transforms.ToTensor()\n",
    "])\n",
    "\n",
    "# Create DataLoaders with help from data_setup.py\n",
    "train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(\n",
    "    train_dir=train_dir,\n",
    "    test_dir=test_dir,\n",
    "    transform=data_transform,\n",
    "    batch_size=BATCH_SIZE\n",
    ")\n",
    "\n",
    "# Create model with help from model_builder.py\n",
    "model = model_builder.TinyVGG(\n",
    "    input_shape=3,\n",
    "    hidden_units=HIDDEN_UNITS,\n",
    "    output_shape=len(class_names)\n",
    ").to(device)\n",
    "\n",
    "# Set loss and optimizer\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(),\n",
    "                             lr=LEARNING_RATE)\n",
    "\n",
    "# Start training with help from engine.py\n",
    "engine.train(model=model,\n",
    "             train_dataloader=train_dataloader,\n",
    "             test_dataloader=test_dataloader,\n",
    "             loss_fn=loss_fn,\n",
    "             optimizer=optimizer,\n",
    "             epochs=NUM_EPOCHS,\n",
    "             device=device)\n",
    "\n",
    "# Save the model with help from utils.py\n",
    "utils.save_model(model=model,\n",
    "                 target_dir=\"models\",\n",
    "                 model_name=\"05_going_modular_script_mode_tinyvgg_model.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "a22d53d4-e54c-419d-819c-15f352088b91",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0%|                                                     | 0/5 [00:00<?, ?it/s]Epoch: 1 | train_loss: 1.0974 | train_acc: 0.4258 | test_loss: 1.1170 | test_acc: 0.2604\n",
      " 20%|█████████                                    | 1/5 [00:00<00:02,  1.55it/s]Epoch: 2 | train_loss: 1.1134 | train_acc: 0.3047 | test_loss: 1.1208 | test_acc: 0.2604\n",
      " 40%|██████████████████                           | 2/5 [00:01<00:01,  2.03it/s]Epoch: 3 | train_loss: 1.1004 | train_acc: 0.3047 | test_loss: 1.0888 | test_acc: 0.4328\n",
      " 60%|███████████████████████████                  | 3/5 [00:01<00:00,  2.35it/s]Epoch: 4 | train_loss: 1.0950 | train_acc: 0.4297 | test_loss: 1.0932 | test_acc: 0.5663\n",
      " 80%|████████████████████████████████████         | 4/5 [00:01<00:00,  2.59it/s]Epoch: 5 | train_loss: 1.0985 | train_acc: 0.3945 | test_loss: 1.1004 | test_acc: 0.2604\n",
      "100%|█████████████████████████████████████████████| 5/5 [00:02<00:00,  2.40it/s]\n",
      "[INFO] Saving model to: models/05_going_modular_script_mode_tinyvgg_model.pth\n"
     ]
    }
   ],
   "source": [
    "!python going_modular/train.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12c972b0-6f42-4151-a807-e659a6b613f3",
   "metadata": {},
   "source": [
    "exercises:\n",
    "1.Turn the code to get the data (from section 1. Get Data above) into a Python script, such as get_data.py.\n",
    "When you run the script using python get_data.py it should check if the data already exists and skip downloading if it does.\n",
    "If the data download is successful, you should be able to access the pizza_steak_sushi images from the data directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "d58588f3-57b1-4342-89fe-71d2be6c52ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing going_modular/get_data.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile going_modular/get_data.py\n",
    "import os\n",
    "import zipfile\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "import requests\n",
    "\n",
    "# Setup path to data folder\n",
    "data_path = Path(\"data/\")\n",
    "image_path = data_path / \"pizza_steak_sushi\"\n",
    "\n",
    "# If the image folder doesn't exist, download it and prepare it... \n",
    "if image_path.is_dir():\n",
    "    print(f\"{image_path} directory exists.\")\n",
    "else:\n",
    "    print(f\"Did not find {image_path} directory, creating one...\")\n",
    "    image_path.mkdir(parents=True, exist_ok=True)\n",
    "    \n",
    "# Download pizza, steak, sushi data\n",
    "with open(data_path / \"pizza_steak_sushi.zip\", \"wb\") as f:\n",
    "    request = requests.get(\"https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip\")\n",
    "    print(\"Downloading pizza, steak, sushi data...\")\n",
    "    f.write(request.content)\n",
    "\n",
    "# Unzip pizza, steak, sushi data\n",
    "with zipfile.ZipFile(data_path / \"pizza_steak_sushi.zip\", \"r\") as zip_ref:\n",
    "    print(\"Unzipping pizza, steak, sushi data...\") \n",
    "    zip_ref.extractall(image_path)\n",
    "\n",
    "# Remove zip file\n",
    "os.remove(data_path / \"pizza_steak_sushi.zip\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "c7f8915d-6b46-4d9a-8d2a-3e4d628dd4b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data/pizza_steak_sushi directory exists.\n",
      "Downloading pizza, steak, sushi data...\n",
      "Unzipping pizza, steak, sushi data...\n"
     ]
    }
   ],
   "source": [
    "!python going_modular/get_data.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b6e7435-fc18-47d1-89ee-ca9fcc904fe8",
   "metadata": {},
   "source": [
    "2.Use Python's argparse module(https://docs.python.org/3/library/argparse.html) to be able to send the train.py custom hyperparameter values for training procedures.\n",
    "Add an argument for using a different:\n",
    "Training/testing directory\n",
    "Learning rate\n",
    "Batch size\n",
    "Number of epochs to train for\n",
    "Number of hidden units in the TinyVGG model\n",
    "Keep the default values for each of the above arguments as what they already are (as in notebook 05).\n",
    "For example, you should be able to run something similar to the following line to train a TinyVGG model with a learning rate of 0.003 and a batch size of 64 for 20 epochs: python train.py --learning_rate 0.003 --batch_size 64 --num_epochs 20.\n",
    "Note: Since train.py leverages the other scripts we created in section 05, such as, model_builder.py, utils.py and engine.py, you'll have to make sure they're available to use too. You can find these in the going_modular folder on the course GitHub."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b61e72e3-d896-4f3d-bdb1-b5b791b083f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing going_modular/train_cli_params.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile going_modular/train_cli_params.py\n",
    "\"\"\"\n",
    "Trains a PyTorch image classification model using device-agnostic code.\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "\n",
    "import torch\n",
    "\n",
    "from torchvision import transforms\n",
    "\n",
    "import data_setup, engine, model_builder, utils\n",
    "\n",
    "import argparse\n",
    "import pathlib\n",
    "\n",
    "parser = argparse.ArgumentParser(description=\"this is the training command line parameters setup explanation\")\n",
    "parser.add_argument('--train_dir', type=pathlib.Path, required=True, help=\"training data directory\")\n",
    "parser.add_argument('--test_dir', type=pathlib.Path, required=True, help=\"test data directory\")\n",
    "parser.add_argument('--learning_rate', type=float, default=0.001, help=\"model training learning rate\")\n",
    "parser.add_argument('--batch_size', type=int, default=32, help=\"model training data batch size\")\n",
    "parser.add_argument('--num_epochs', type=int, default=5, help=\"model training epochs number\")\n",
    "parser.add_argument('--hidden_units', type=int, default=10, help=\"training model structure hidden units number\")\n",
    "args = parser.parse_args()\n",
    "\n",
    "# Setup hyperparameters\n",
    "NUM_EPOCHS = args.num_epochs\n",
    "BATCH_SIZE = args.batch_size\n",
    "HIDDEN_UNITS = args.hidden_units\n",
    "LEARNING_RATE = args.learning_rate\n",
    "\n",
    "# Setup directories\n",
    "train_dir = args.train_dir\n",
    "test_dir = args.test_dir\n",
    "\n",
    "# Setup target device\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# Create transforms\n",
    "data_transform = transforms.Compose([\n",
    "  transforms.Resize((64, 64)),\n",
    "  transforms.ToTensor()\n",
    "])\n",
    "\n",
    "# Create DataLoaders with help from data_setup.py\n",
    "train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(\n",
    "    train_dir=train_dir,\n",
    "    test_dir=test_dir,\n",
    "    transform=data_transform,\n",
    "    batch_size=BATCH_SIZE\n",
    ")\n",
    "\n",
    "# Create model with help from model_builder.py\n",
    "model = model_builder.TinyVGG(\n",
    "    input_shape=3,\n",
    "    hidden_units=HIDDEN_UNITS,\n",
    "    output_shape=len(class_names)\n",
    ").to(device)\n",
    "\n",
    "# Set loss and optimizer\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(),\n",
    "                             lr=LEARNING_RATE)\n",
    "\n",
    "# Start training with help from engine.py\n",
    "engine.train(model=model,\n",
    "             train_dataloader=train_dataloader,\n",
    "             test_dataloader=test_dataloader,\n",
    "             loss_fn=loss_fn,\n",
    "             optimizer=optimizer,\n",
    "             epochs=NUM_EPOCHS,\n",
    "             device=device)\n",
    "\n",
    "# Save the model with help from utils.py\n",
    "utils.save_model(model=model,\n",
    "                 target_dir=\"models\",\n",
    "                 model_name=\"05_going_modular_script_mode_tinyvgg_model_with_cli_params.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "3de92a78-9525-4ce3-831a-a8dde867e92a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0%|                                                    | 0/20 [00:00<?, ?it/s]Epoch: 1 | train_loss: 1.1251 | train_acc: 0.2569 | test_loss: 1.1538 | test_acc: 0.1484\n",
      "  5%|██▏                                         | 1/20 [00:00<00:17,  1.09it/s]Epoch: 2 | train_loss: 1.1086 | train_acc: 0.3223 | test_loss: 1.0950 | test_acc: 0.3771\n",
      " 10%|████▍                                       | 2/20 [00:01<00:11,  1.62it/s]Epoch: 3 | train_loss: 1.0922 | train_acc: 0.3414 | test_loss: 1.0975 | test_acc: 0.2031\n",
      " 15%|██████▌                                     | 3/20 [00:01<00:08,  1.92it/s]Epoch: 4 | train_loss: 1.0680 | train_acc: 0.4576 | test_loss: 1.0447 | test_acc: 0.4915\n",
      " 20%|████████▊                                   | 4/20 [00:02<00:07,  2.11it/s]Epoch: 5 | train_loss: 0.9798 | train_acc: 0.5575 | test_loss: 1.1238 | test_acc: 0.2955\n",
      " 25%|███████████                                 | 5/20 [00:02<00:06,  2.23it/s]Epoch: 6 | train_loss: 0.8863 | train_acc: 0.5966 | test_loss: 0.8573 | test_acc: 0.6499\n",
      " 30%|█████████████▏                              | 6/20 [00:03<00:06,  2.12it/s]Epoch: 7 | train_loss: 0.9456 | train_acc: 0.5228 | test_loss: 1.0896 | test_acc: 0.3629\n",
      " 35%|███████████████▍                            | 7/20 [00:03<00:05,  2.22it/s]Epoch: 8 | train_loss: 0.8523 | train_acc: 0.6232 | test_loss: 1.2337 | test_acc: 0.2266\n",
      " 40%|█████████████████▌                          | 8/20 [00:03<00:05,  2.29it/s]Epoch: 9 | train_loss: 0.9156 | train_acc: 0.5507 | test_loss: 1.1473 | test_acc: 0.2642\n",
      " 45%|███████████████████▊                        | 9/20 [00:04<00:04,  2.33it/s]Epoch: 10 | train_loss: 0.8833 | train_acc: 0.6085 | test_loss: 1.0241 | test_acc: 0.4460\n",
      " 50%|█████████████████████▌                     | 10/20 [00:04<00:04,  2.34it/s]Epoch: 11 | train_loss: 0.8171 | train_acc: 0.6207 | test_loss: 0.9780 | test_acc: 0.4084\n",
      " 55%|███████████████████████▋                   | 11/20 [00:05<00:03,  2.37it/s]Epoch: 12 | train_loss: 0.8164 | train_acc: 0.6625 | test_loss: 1.0274 | test_acc: 0.5149\n",
      " 60%|█████████████████████████▊                 | 12/20 [00:05<00:03,  2.36it/s]Epoch: 13 | train_loss: 0.7829 | train_acc: 0.6818 | test_loss: 1.1787 | test_acc: 0.4318\n",
      " 65%|███████████████████████████▉               | 13/20 [00:05<00:02,  2.38it/s]Epoch: 14 | train_loss: 0.7429 | train_acc: 0.6970 | test_loss: 1.1116 | test_acc: 0.4162\n",
      " 70%|██████████████████████████████             | 14/20 [00:06<00:02,  2.38it/s]Epoch: 15 | train_loss: 0.7282 | train_acc: 0.6781 | test_loss: 1.0732 | test_acc: 0.3928\n",
      " 75%|████████████████████████████████▎          | 15/20 [00:06<00:02,  2.39it/s]Epoch: 16 | train_loss: 0.7233 | train_acc: 0.6896 | test_loss: 1.0707 | test_acc: 0.4006\n",
      " 80%|██████████████████████████████████▍        | 16/20 [00:07<00:01,  2.38it/s]Epoch: 17 | train_loss: 0.6919 | train_acc: 0.6864 | test_loss: 1.1090 | test_acc: 0.3864\n",
      " 85%|████████████████████████████████████▌      | 17/20 [00:07<00:01,  2.40it/s]Epoch: 18 | train_loss: 0.6632 | train_acc: 0.7135 | test_loss: 1.1046 | test_acc: 0.3786\n",
      " 90%|██████████████████████████████████████▋    | 18/20 [00:08<00:00,  2.42it/s]Epoch: 19 | train_loss: 0.6500 | train_acc: 0.7414 | test_loss: 1.2417 | test_acc: 0.3707\n",
      " 95%|████████████████████████████████████████▊  | 19/20 [00:08<00:00,  2.43it/s]Epoch: 20 | train_loss: 0.5863 | train_acc: 0.7673 | test_loss: 1.4227 | test_acc: 0.3771\n",
      "100%|███████████████████████████████████████████| 20/20 [00:08<00:00,  2.26it/s]\n",
      "[INFO] Saving model to: models/05_going_modular_script_mode_tinyvgg_model_with_cli_params.pth\n"
     ]
    }
   ],
   "source": [
    "!python going_modular/train_cli_params.py  --train_dir \"data/pizza_steak_sushi/train\" --test_dir \"data/pizza_steak_sushi/test\" --learning_rate 0.003 --batch_size 64 --num_epochs 20 --hidden_units 20"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d97df2af-ee42-41b5-902e-c55513a8d0e4",
   "metadata": {},
   "source": [
    "3.Create a script to predict (such as predict.py) on a target image given a file path with a saved model.\n",
    "For example, you should be able to run the command python predict.py some_image.jpeg and have a trained PyTorch model predict on the image and return its prediction.\n",
    "To see example prediction code, check out the predicting on a custom image section in notebook 04(https://www.learnpytorch.io/04_pytorch_custom_datasets/#113-putting-custom-image-prediction-together-building-a-function).\n",
    "You may also have to write code to load in a trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "dd8a545b-08d8-4869-a971-e0f38aaecd90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing going_modular/predict.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile going_modular/predict.py\n",
    "\"\"\"\n",
    "Predict a PyTorch image using device-agnostic code.\n",
    "\"\"\"\n",
    "import argparse\n",
    "import pathlib\n",
    "from pathlib import Path\n",
    "import model_builder\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "import torchvision\n",
    "from typing import List\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# Setup path to data folder\n",
    "data_path = Path(\"data/\")\n",
    "image_path = data_path / \"pizza_steak_sushi\"\n",
    "train_dir = image_path / \"train\"\n",
    "\n",
    "parser = argparse.ArgumentParser(description=\"this is the inference prediction command line parameters explanation\")\n",
    "parser.add_argument('--predict_image', type=pathlib.Path, required=True, help=\"image prediction filepath\")\n",
    "args = parser.parse_args()\n",
    "\n",
    "custom_image_path = data_path / args.predict_image\n",
    "\n",
    "MODEL_PATH = Path(\"models\")\n",
    "MODEL_NAME = \"05_going_modular_script_mode_tinyvgg_model_with_cli_params.pth\"\n",
    "MODEL_SAVE_PATH = MODEL_PATH / MODEL_NAME\n",
    "\n",
    "\n",
    "loaded_model = model_builder.TinyVGG(input_shape=3, \n",
    "                                    hidden_units=20, # try changing this to 128 and seeing what happens \n",
    "                                    output_shape=3) \n",
    "\n",
    "# Load in the saved state_dict()\n",
    "loaded_model.load_state_dict(torch.load(f=MODEL_SAVE_PATH))\n",
    "\n",
    "# Send model to GPU\n",
    "loaded_model = loaded_model.to(device)\n",
    "\n",
    "# Create transform pipleine to resize image\n",
    "custom_image_transform = transforms.Compose([\n",
    "    transforms.Resize((64, 64)),\n",
    "])\n",
    "\n",
    "\n",
    "def pred_image(model: torch.nn.Module, \n",
    "                        image_path: str, \n",
    "                        class_names: List[str] = None, \n",
    "                        transform=None,\n",
    "                        device: torch.device = device):\n",
    "    \"\"\"Makes a prediction on a target image and plots the image with its prediction.\"\"\"\n",
    "    \n",
    "    # 1. Load in image and convert the tensor values to float32\n",
    "    target_image = torchvision.io.read_image(str(image_path)).type(torch.float32)\n",
    "    \n",
    "    # 2. Divide the image pixel values by 255 to get them between [0, 1]\n",
    "    target_image = target_image / 255. \n",
    "    \n",
    "    # 3. Transform if necessary\n",
    "    if transform:\n",
    "        target_image = transform(target_image)\n",
    "    \n",
    "    # 4. Make sure the model is on the target device\n",
    "    model.to(device)\n",
    "    \n",
    "    # 5. Turn on model evaluation mode and inference mode\n",
    "    model.eval()\n",
    "    with torch.inference_mode():\n",
    "        # Add an extra dimension to the image\n",
    "        target_image = target_image.unsqueeze(dim=0)\n",
    "    \n",
    "        # Make a prediction on image with an extra dimension and send it to the target device\n",
    "        target_image_pred = model(target_image.to(device))\n",
    "        \n",
    "    # 6. Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)\n",
    "    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)\n",
    "\n",
    "    # 7. Convert prediction probabilities -> prediction labels\n",
    "    target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)\n",
    "    \n",
    "    # 8. prediction probability\n",
    "\n",
    "    if class_names:\n",
    "        result = f\"Pred: {class_names[target_image_pred_label.cpu()]} | Prob: {target_image_pred_probs.max().cpu():.3f}\"\n",
    "    else: \n",
    "        result = f\"Pred: {target_image_pred_label} | Prob: {target_image_pred_probs.max().cpu():.3f}\"\n",
    "    return result\n",
    "\n",
    "# Use ImageFolder to create dataset(s)\n",
    "from torchvision import datasets\n",
    "train_data = datasets.ImageFolder(root=train_dir, # target folder of images\n",
    "                                  transform=custom_image_transform, # transforms to perform on data (images)\n",
    "                                  target_transform=None) # transforms to perform on labels (if necessary)\n",
    "# Get class names as a list\n",
    "class_names = train_data.classes\n",
    "\n",
    "# Pred on our custom image\n",
    "pred_result = pred_image(model=loaded_model,\n",
    "                    image_path=custom_image_path,\n",
    "                    class_names=class_names,\n",
    "                    transform=custom_image_transform,\n",
    "                    device=device)\n",
    "print(pred_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "b56fc17a-02a3-4de8-8996-93eee4b125f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/qm/agui/miniconda/envs/pytorch/lib/python3.11/site-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).\n",
      "  warnings.warn(\n",
      "Pred: steak | Prob: 0.819\n"
     ]
    }
   ],
   "source": [
    "!python going_modular/predict.py --predict_image \"04-pizza-dad.jpeg\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95f2604b-6d42-496c-b846-f1d73dd76c2b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}