05_IFCB_CNN_condensed_cellbased_test_train.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOkyeFDjh0y/+QGxhQeLIrp",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "1a98b23709554441bfc01390cf5a7444": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_c8b636391cf04a2b9c3368cdadd4c340",
              "IPY_MODEL_e4efea977f6f433890b91a56115e97be",
              "IPY_MODEL_f9d8e9f1d5414675aadeaa1acfbd1bab"
            ],
            "layout": "IPY_MODEL_d692a061058c46619fdd172ba2f94d1e"
          }
        },
        "c8b636391cf04a2b9c3368cdadd4c340": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a0e8d6415d674dadb4d3e6bc53f0618a",
            "placeholder": "​",
            "style": "IPY_MODEL_923c5194518943f69f2cee79c81b8fd1",
            "value": "100%"
          }
        },
        "e4efea977f6f433890b91a56115e97be": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ae4755011ca0442c906d4c44c95aeb20",
            "max": 5,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_713a77b40bca45989195be04df11fcea",
            "value": 5
          }
        },
        "f9d8e9f1d5414675aadeaa1acfbd1bab": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5ff321e4ec484e219397447506ff42cd",
            "placeholder": "​",
            "style": "IPY_MODEL_96fa7920f7344336ae87f6f5da5e7303",
            "value": " 5/5 [00:16&lt;00:00,  3.18s/it]"
          }
        },
        "d692a061058c46619fdd172ba2f94d1e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a0e8d6415d674dadb4d3e6bc53f0618a": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "923c5194518943f69f2cee79c81b8fd1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "ae4755011ca0442c906d4c44c95aeb20": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "713a77b40bca45989195be04df11fcea": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "5ff321e4ec484e219397447506ff42cd": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "96fa7920f7344336ae87f6f5da5e7303": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/maritnorli/IFCB_CNN_Classify/blob/main/05_IFCB_CNN_condensed_cellbased_test_train.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 0. Running a notebook in cell mode\n",
        "As discussed, we're going to be running this notebook normally.\n",
        "\n",
        "One cell at a time.\n",
        "\n",
        "The code is from notebook 04, however, it has been condensed down to its core functionality."
      ],
      "metadata": {
        "id": "3iHTyE5CJhev"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "1. Get data\n",
        "We're going to start by downloading the same data we used in notebook 04."
      ],
      "metadata": {
        "id": "gVGPfGAGJzqa"
      }
    },
    {
      "cell_type": "code",
      "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 / \"IFCB_test_train\"\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 / \"IFCB_test_train.zip\", \"wb\") as f:\n",
        "    request = requests.get(\"https://github.com/maritnorli/IFCB_CNN_Classify/raw/main/data/IFCB_test_train.zip\")\n",
        "    print(\"Downloading IFCB data...\")\n",
        "    f.write(request.content)\n",
        "\n",
        "# Unzip pizza, steak, sushi data\n",
        "with zipfile.ZipFile(data_path / \"IFCB_test_train.zip\", \"r\") as zip_ref:\n",
        "    print(\"Unzipping IFCB_test_train data...\")\n",
        "    zip_ref.extractall(image_path)\n",
        "\n",
        "# Remove zip file\n",
        "os.remove(data_path / \"IFCB_test_train.zip\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xCl8frgxJ0YG",
        "outputId": "4406e1ff-abef-408e-d991-8db367f37c77"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Did not find data/IFCB_test_train directory, creating one...\n",
            "Downloading IFCB data...\n",
            "Unzipping IFCB_test_train data...\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Setup train and testing paths\n",
        "train_dir = image_path / \"Train\"\n",
        "test_dir = image_path / \"Test\"\n",
        "\n",
        "train_dir, test_dir"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eYZERVCxJ0UB",
        "outputId": "da72f71d-95cb-4b60-fbab-cdec8b51c0a2"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(PosixPath('data/IFCB_test_train/Train'),\n",
              " PosixPath('data/IFCB_test_train/Test'))"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2. Create Datasets and DataLoaders\n",
        "\n",
        "Now we'll turn the image dataset into PyTorch `Dataset`'s and `DataLoader`'s."
      ],
      "metadata": {
        "id": "onQbKQp3J0Q5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torchvision import datasets, transforms\n",
        "\n",
        "# Create simple transform\n",
        "data_transform = transforms.Compose([\n",
        "    transforms.Grayscale(),\n",
        "    transforms.Resize(size=(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}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ox4vgIlqJ0Ow",
        "outputId": "77b749d8-b3a9-4de4-dedd-9ce0fe15ea07"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train data:\n",
            "Dataset ImageFolder\n",
            "    Number of datapoints: 210\n",
            "    Root location: data/IFCB_test_train/Train\n",
            "    StandardTransform\n",
            "Transform: Compose(\n",
            "               Grayscale(num_output_channels=1)\n",
            "               Resize(size=(64, 64), interpolation=bilinear, max_size=None, antialias=True)\n",
            "               ToTensor()\n",
            "           )\n",
            "Test data:\n",
            "Dataset ImageFolder\n",
            "    Number of datapoints: 42\n",
            "    Root location: data/IFCB_test_train/Test\n",
            "    StandardTransform\n",
            "Transform: Compose(\n",
            "               Grayscale(num_output_channels=1)\n",
            "               Resize(size=(64, 64), interpolation=bilinear, max_size=None, antialias=True)\n",
            "               ToTensor()\n",
            "           )\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Get class names as a list\n",
        "class_names = train_data.classes\n",
        "class_names"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VlFj6w2FJ0Lq",
        "outputId": "4bb584fb-002e-4431-f50f-c8bbb119630e"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['Chaetoceros_decipiens_118',\n",
              " 'Guinardia_delicatula_095',\n",
              " 'Tripos_muelleri_008']"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Can also get class names as a dict\n",
        "class_dict = train_data.class_to_idx\n",
        "class_dict"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NUn5F8DnJ0Ej",
        "outputId": "4f353d97-ed2b-4206-9aba-1295e091898b"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'Chaetoceros_decipiens_118': 0,\n",
              " 'Guinardia_delicatula_095': 1,\n",
              " 'Tripos_muelleri_008': 2}"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Check the lengths\n",
        "len(train_data), len(test_data)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bUAl3Yy9LQJT",
        "outputId": "177b0d02-f482-46c6-e303-3ea4f32124f7"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(210, 42)"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Turn train and test Datasets into DataLoaders\n",
        "from torch.utils.data import DataLoader\n",
        "train_dataloader = DataLoader(dataset=train_data,\n",
        "                              batch_size=1, # how many samples per batch?\n",
        "                              num_workers=1, # 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=1,\n",
        "                             num_workers=1,\n",
        "                             shuffle=False) # don't usually need to shuffle testing data\n",
        "\n",
        "train_dataloader, test_dataloader"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lm4gYMxHLQFk",
        "outputId": "14c7e8af-1ffc-493e-c0aa-23d6fea3da28"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(<torch.utils.data.dataloader.DataLoader at 0x7fabb55cee30>,\n",
              " <torch.utils.data.dataloader.DataLoader at 0x7fabb55cea40>)"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "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}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QHp2VC9RLQCT",
        "outputId": "4615db37-8179-44f8-87df-70377403bcfe"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Image shape: torch.Size([1, 1, 64, 64]) -> [batch_size, color_channels, height, width]\n",
            "Label shape: torch.Size([1])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 3. Making a model (TinyVGG)\n",
        "We're going to use the same model we used in notebook 04: TinyVGG from the CNN Explainer website.\n",
        "\n",
        "The only change here from notebook 04 is that a docstring has been added using Google's Style Guide for Python."
      ],
      "metadata": {
        "id": "4S9cChVcLP-r"
      }
    },
    {
      "cell_type": "code",
      "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, # how big is the square that's going over the image?\n",
        "                    stride=1, # default\n",
        "                    padding=0), # options = \"valid\" (no padding) or \"same\" (output has same shape as input) or int for specific number\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) # default stride value is same as kernel_size\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"
      ],
      "metadata": {
        "id": "fNUTlcdyLP7J"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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=1, # number of color channels (3 for RGB)\n",
        "                  hidden_units=10,\n",
        "                  output_shape=len(train_data.classes)).to(device)\n",
        "model_0"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fndI0flUL--p",
        "outputId": "ead54166-9e3a-4367-a162-3cf9eebc97cd"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TinyVGG(\n",
              "  (conv_block_1): Sequential(\n",
              "    (0): Conv2d(1, 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",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "To test our model let's do a single forward pass (pass a sample batch from the training set through our model)."
      ],
      "metadata": {
        "id": "yTzuIaECL_1W"
      }
    },
    {
      "cell_type": "code",
      "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}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f88qjEnfL_xU",
        "outputId": "c116a571-5c08-49d1-b1a4-17773084b5fb"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Single image shape: torch.Size([1, 1, 64, 64])\n",
            "\n",
            "Output logits:\n",
            "tensor([[-0.0139,  0.0263,  0.0175]])\n",
            "\n",
            "Output prediction probabilities:\n",
            "tensor([[0.3254, 0.3388, 0.3358]])\n",
            "\n",
            "Output prediction label:\n",
            "tensor([1])\n",
            "\n",
            "Actual label:\n",
            "1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 4. Creating `train_step()` and `test_step()` functions and `train()` to combine them  \n",
        "Rather than writing them again, we can reuse the train_step() and test_step() functions from notebook 04.\n",
        "\n",
        "The same goes for the train() function we created.\n",
        "\n",
        "The only difference here is that these functions have had docstrings added to them in Google's Python Functions and Methods Style Guide.\n",
        "\n",
        "Let's start by making train_step()."
      ],
      "metadata": {
        "id": "pVKhEcvVL_ui"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "EsX7_bIaL_rY"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we'll do test_step().\n",
        "\n"
      ],
      "metadata": {
        "id": "YRiHOCSKL_oV"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "-AKgvfJ0L_lO"
      },
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And we'll combine train_step() and test_step() into train()."
      ],
      "metadata": {
        "id": "UbJG1oe9L_iK"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "me9h6kXvq7sT"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 5. Creating a function to save the model\n",
        "\n",
        "Let's setup a function to save our model to a directory."
      ],
      "metadata": {
        "id": "TUCBwMuYrD4G"
      }
    },
    {
      "cell_type": "code",
      "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)"
      ],
      "metadata": {
        "id": "68ivpt0qrHwO"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 6. Train, evaluate and save the model\n",
        "\n",
        "Let's leverage the functions we've got above to train, test and save a model to file.\n"
      ],
      "metadata": {
        "id": "O30Pj6OprObK"
      }
    },
    {
      "cell_type": "code",
      "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=1, # 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\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 193,
          "referenced_widgets": [
            "1a98b23709554441bfc01390cf5a7444",
            "c8b636391cf04a2b9c3368cdadd4c340",
            "e4efea977f6f433890b91a56115e97be",
            "f9d8e9f1d5414675aadeaa1acfbd1bab",
            "d692a061058c46619fdd172ba2f94d1e",
            "a0e8d6415d674dadb4d3e6bc53f0618a",
            "923c5194518943f69f2cee79c81b8fd1",
            "ae4755011ca0442c906d4c44c95aeb20",
            "713a77b40bca45989195be04df11fcea",
            "5ff321e4ec484e219397447506ff42cd",
            "96fa7920f7344336ae87f6f5da5e7303"
          ]
        },
        "id": "o7JRbUqHrOUZ",
        "outputId": "f428404d-117c-4d67-b0bb-3e3d62d5e871"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/5 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "1a98b23709554441bfc01390cf5a7444"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 1 | train_loss: 1.1066 | train_acc: 0.3286 | test_loss: 1.0977 | test_acc: 0.3333\n",
            "Epoch: 2 | train_loss: 1.0864 | train_acc: 0.3143 | test_loss: 1.0394 | test_acc: 0.4524\n",
            "Epoch: 3 | train_loss: 0.9169 | train_acc: 0.5952 | test_loss: 0.5786 | test_acc: 0.7143\n",
            "Epoch: 4 | train_loss: 0.4812 | train_acc: 0.8333 | test_loss: 0.3811 | test_acc: 0.8333\n",
            "Epoch: 5 | train_loss: 0.4374 | train_acc: 0.8476 | test_loss: 0.3770 | test_acc: 0.8333\n",
            "[INFO] Total training time: 16.179 seconds\n",
            "[INFO] Saving model to: models/05_going_modular_cell_mode_tinyvgg_model.pth\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We finish with a saved image classification model at `models/05_going_modular_cell_mode_tinyvgg_model.pth`.\n",
        "\n",
        "The code continues in [05. Going Modular: Part 2 (script mode)](https://github.com/mrdbourke/pytorch-deep-learning/blob/main/going_modular/05_pytorch_going_modular_script_mode.ipynb)."
      ],
      "metadata": {
        "id": "vKTgQ9R7rOOp"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "cSOsQnN-rOFN"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}