{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d6112ee2-1800-4086-b7c4-cc89ff45a06a",
   "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 requests\n",
    "import zipfile\n",
    "from pathlib import Path\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": 2,
   "id": "4ff2a492-fd34-41e9-9a82-ab63ff7d3215",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing 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 torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\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, # don't need to shuffle test data\n",
    "      num_workers=num_workers,\n",
    "      pin_memory=True,\n",
    "  )\n",
    "\n",
    "  return train_dataloader, test_dataloader, class_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d8a44d9a-b71f-4835-b1f5-44ba1203d875",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing 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",
    "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": 4,
   "id": "53d0f879-f508-4484-8d9b-2bcb97484eff",
   "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",
    "import torch\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "from typing import Dict, List, 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\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]:\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": 5,
   "id": "fceb758b-a949-4bf3-b11f-68d37563f24a",
   "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",
    "import torch\n",
    "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": null,
   "id": "a8b901f4-8847-4308-af41-afb0de9b8cd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile going_modular/train.py\n",
    "\"\"\"\n",
    "Trains a PyTorch image classification model using device-agnostic code.\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "import torch\n",
    "import data_setup, engine, model_builder, utils\n",
    "\n",
    "from torchvision import transforms\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 = \"C:/Users/User/Desktop/Pytorch/pytorchPractice/data/pizza_steak_sushi/train\"\n",
    "test_dir = \"C:/Users/User/Desktop/Pytorch/pytorchPractice/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=\"C:/Users/User/Desktop/Pytorch/pytorchPractice/models\",\n",
    "                 model_name=\"05_going_modular_script_mode_tinyvgg_model.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20f6887d-b3b4-484f-9d74-6823a246d7e3",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}