Skip to content

Instantly share code, notes, and snippets.

@dnkm

dnkm/ai-powered-air-quality-prediction-system.ipynb

Created June 24, 2026 22:03
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save dnkm/247d5e3d567e660e50b5f1ed400b3943 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save dnkm/247d5e3d567e660e50b5f1ed400b3943 to your computer and use it in GitHub Desktop.
Download ZIP
AI powered Air Quality Prediction System
Display the source blob
Display the rendered blob
Raw
ai-powered-air-quality-prediction-system.ipynb
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/dnkm/247d5e3d567e660e50b5f1ed400b3943/ai-powered-air-quality-prediction-system.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"source": [
"# **AI Powered Air Quality Prediction System**"
],
"metadata": {
"id": "6N7XTwulh_xI"
}
},
{
"cell_type": "markdown",
"metadata": {
"id": "Ifzv2FxT2XnK"
},
"source": [
"**Joowon Kang**"
]
},
{
"cell_type": "markdown",
"source": [
"**Instructions to load data**\n",
"\n",
"1. Download the dataset zip file\n",
"2. Upload the zip file onto your Google Drive\n",
"3. Run the code below!"
],
"metadata": {
"id": "o-mdb7MAirb6"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "pq4N5WKSy4Ex"
},
"outputs": [],
"source": [
"import os\n",
"import cv2\n",
"import math\n",
"import torch\n",
"import numpy as np\n",
"import pandas as pd\n",
"import torch.nn as nn\n",
"import matplotlib.pyplot as plt"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "mwhH3T-eswTH"
},
"outputs": [],
"source": [
"from PIL import Image\n",
"from google.colab import drive\n",
"from torchvision import models\n",
"from torchvision import transforms\n",
"from collections import OrderedDict\n",
"from torch.utils.data import Dataset, DataLoader"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 293
},
"id": "4C5jzcnqtNMI",
"outputId": "3f14fec9-b64f-4cf5-8377-8a8525acc9b5"
},
"outputs": [
{
"output_type": "error",
"ename": "MessageError",
"evalue": "Error: credential propagation was unsuccessful",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mMessageError\u001b[0m Traceback (most recent call last)",
"\u001b[0;32m/tmp/ipykernel_12805/2985678062.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mdrive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmount\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'/content/drive'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
"\u001b[0;32m/usr/local/lib/python3.12/dist-packages/google/colab/drive.py\u001b[0m in \u001b[0;36mmount\u001b[0;34m(mountpoint, force_remount, timeout_ms, readonly)\u001b[0m\n\u001b[1;32m 95\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mmount\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmountpoint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mforce_remount\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtimeout_ms\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m120000\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreadonly\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 96\u001b[0m \u001b[0;34m\"\"\"Mount your Google Drive at the specified mountpoint path.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 97\u001b[0;31m return _mount(\n\u001b[0m\u001b[1;32m 98\u001b[0m \u001b[0mmountpoint\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 99\u001b[0m \u001b[0mforce_remount\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mforce_remount\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.12/dist-packages/google/colab/drive.py\u001b[0m in \u001b[0;36m_mount\u001b[0;34m(mountpoint, force_remount, timeout_ms, ephemeral, readonly)\u001b[0m\n\u001b[1;32m 132\u001b[0m )\n\u001b[1;32m 133\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mephemeral\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 134\u001b[0;31m _message.blocking_request(\n\u001b[0m\u001b[1;32m 135\u001b[0m \u001b[0;34m'request_auth'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 136\u001b[0m \u001b[0mrequest\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0;34m'authType'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'dfs_ephemeral'\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.12/dist-packages/google/colab/_message.py\u001b[0m in \u001b[0;36mblocking_request\u001b[0;34m(request_type, request, timeout_sec, parent)\u001b[0m\n\u001b[1;32m 174\u001b[0m \u001b[0mrequest_type\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrequest\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparent\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mparent\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexpect_reply\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 175\u001b[0m )\n\u001b[0;32m--> 176\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mread_reply_from_input\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrequest_id\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtimeout_sec\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
"\u001b[0;32m/usr/local/lib/python3.12/dist-packages/google/colab/_message.py\u001b[0m in \u001b[0;36mread_reply_from_input\u001b[0;34m(message_id, timeout_sec)\u001b[0m\n\u001b[1;32m 101\u001b[0m ):\n\u001b[1;32m 102\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;34m'error'\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mreply\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 103\u001b[0;31m \u001b[0;32mraise\u001b[0m \u001b[0mMessageError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mreply\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'error'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 104\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mreply\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'data'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 105\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;31mMessageError\u001b[0m: Error: credential propagation was unsuccessful"
]
}
],
"source": [
"drive.mount('/content/drive')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "XWx9QylCPpVe"
},
"outputs": [],
"source": [
"!unzip -q /content/drive/MyDrive/Dataset.zip -d /content/DatasetFolder"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "kHTn3t5GtWUo"
},
"outputs": [],
"source": [
"IMAGE_DIR = \"/content/DatasetFolder/Dataset/Air_project\"\n",
"SHEET_PATH = \"/content/DatasetFolder/Dataset/Data_Recording_Sheet.xlsx\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Sljxb5N17VZh"
},
"outputs": [],
"source": [
"sheet = pd.read_excel(SHEET_PATH)"
]
},
{
"cell_type": "markdown",
"source": [
"**Splitting Data into Train, Test, and Overfitting Sets**"
],
"metadata": {
"id": "NkSBxjHyng9n"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2TlSWItnMNnJ"
},
"outputs": [],
"source": [
"# import pandas as pd\n",
"# from sklearn.model_selection import train_test_split\n",
"# import os\n",
"\n",
"# # Define the file path and desired ratios\n",
"# # file_path = SHEET_PATH\n",
"# train_ratio = 0.8\n",
"# test_ratio = 0.1\n",
"# validation_ratio = 0.1\n",
"# seed = 42 # for reproducibility\n",
"\n",
"# # 1. Read the Excel file\n",
"# df = pd.read_excel(SHEET_PATH)\n",
"\n",
"# sampled_df = df.sample(n=200, random_state=None)\n",
"\n",
"# output_dir = '/content/drive/MyDrive/data_splits'\n",
"# os.makedirs(output_dir, exist_ok=True)\n",
"# sampled_df.to_excel(os.path.join(output_dir, 'overfit_set.xlsx'), index=False)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "pBeGN-j7yWGY"
},
"outputs": [],
"source": [
"# import pandas as pd\n",
"# from sklearn.model_selection import train_test_split\n",
"# import os\n",
"\n",
"# # Define the file path and desired ratios\n",
"# # file_path = SHEET_PATH\n",
"# train_ratio = 0.8\n",
"# test_ratio = 0.1\n",
"# validation_ratio = 0.1\n",
"# seed = 42 # for reproducibility\n",
"\n",
"# # 1. Read the Excel file\n",
"# df = pd.read_excel(SHEET_PATH)\n",
"\n",
"# # 2. Split into main (train+val) and test sets (e.g., 90% main, 10% test)\n",
"# # The test_size for the first split is the final desired test ratio\n",
"# df_main, df_test = train_test_split(df, test_size=test_ratio, random_state=seed)\n",
"\n",
"# # 3. Split the main set into train and validation sets\n",
"# # The validation set size must be adjusted relative to the remaining main set\n",
"# # New validation ratio = actual_validation_ratio / (actual_train_ratio + actual_validation_ratio)\n",
"# val_ratio_adjusted = validation_ratio / (train_ratio + validation_ratio)\n",
"\n",
"# df_train, df_val = train_test_split(df_main, test_size=val_ratio_adjusted, random_state=seed)\n",
"\n",
"# # 4. Save the splits to new Excel files\n",
"# output_dir = '/content/drive/MyDrive/data_splits'\n",
"# os.makedirs(output_dir, exist_ok=True)\n",
"\n",
"# df_train.to_excel(os.path.join(output_dir, 'train_set.xlsx'), index=False)\n",
"# df_val.to_excel(os.path.join(output_dir, 'val_set.xlsx'), index=False)\n",
"# df_test.to_excel(os.path.join(output_dir, 'test_set.xlsx'), index=False)\n",
"\n",
"# print(f\"\\nData successfully split and saved in the '{output_dir}' directory.\")"
]
},
{
"cell_type": "markdown",
"source": [
"**Converting Air Quality Values into Air Quality Labels using US EPA Offical Conversion Tool: https://1drv.ms/w/c/ba0921185736d6f8/EfW2J4_I9KFBjxfDGaXzo5EBdlLlNNhN_tXeTBXPUpAMGQ?e=lEedJ7**"
],
"metadata": {
"id": "57cdO_R-nuDj"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "eybYimTtyrKu"
},
"outputs": [],
"source": [
"label_columns = [\"PM2.5 µg/m³\", \"PM10 µg/m³\", \"CO ppm\", \" NO₂ ppm\", \" SO₂ ppm\", \"O₃ ppm\"]\n",
"air_measures = {row[\"Name\"]: row[label_columns].tolist() for _, row in sheet.iterrows()}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rxTbkta68hfK"
},
"outputs": [],
"source": [
"BREAKPOINTS = {\n",
" \"PM25\": [ (0.0, 9.0, 0, 50),\n",
" (9.1, 35.4, 51, 100),\n",
" (35.5, 55.4, 101, 150),\n",
" (55.5, 125.4, 151, 200),\n",
" (125.5, 225.4, 201, 300),\n",
" (225.5, math.inf, 301, math.inf), ],\n",
" \"PM10\": [ (0, 54, 0, 50),\n",
" (55, 154, 51, 100),\n",
" (155, 254, 101, 150),\n",
" (255, 354, 151, 200),\n",
" (355, 424, 201, 300),\n",
" (425, math.inf, 301, math.inf), ],\n",
" \"O3\": [ (0.000, 0.054, 0, 50),\n",
" (0.055, 0.070, 51, 100),\n",
" (0.071, 0.085, 101, 150),\n",
" (0.086, 0.105, 151, 200),\n",
" (0.106, 0.200, 201, 300),\n",
" (0.201, math.inf, 301, math.inf)],\n",
" \"CO\": [ (0.0, 4.4, 0, 50),\n",
" (4.5, 9.4, 51, 100),\n",
" (9.5, 12.4, 101, 150),\n",
" (12.5, 15.4, 151, 200),\n",
" (15.5, 30.4, 201, 300),\n",
" (30.5, math.inf, 301, math.inf), ],\n",
" \"SO2\": [ (0, 35, 0, 50),\n",
" (36, 75, 51, 100),\n",
" (76, 185, 101, 150),\n",
" (186, 304, 151, 200),\n",
" (305, 604, 201, 300),\n",
" (605, math.inf, 301, math.inf), ],\n",
" \"NO2\": [ (0, 53, 0, 50),\n",
" (54, 100, 51, 100),\n",
" (101, 360, 101, 150),\n",
" (361, 649, 151, 200),\n",
" (650, 1249, 201, 300),\n",
" (1250, math.inf, 301, math.inf), ], }"
]
},
{
"cell_type": "markdown",
"source": [
"**PyTorch Dataloader for Dataset**"
],
"metadata": {
"id": "oWHtW5QMo46u"
}
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "YWMbCYrhy5iW"
},
"outputs": [],
"source": [
"class AirQualityDataset(Dataset):\n",
" def __init__(self, folder, split, transform=None):\n",
" file_path = f'/content/DatasetFolder/Dataset/data_splits/{split}_set.xlsx'\n",
" sheet = pd.read_excel(file_path)\n",
" self.folder = folder\n",
" self.transform = transform\n",
" self.label_columns = [\"PM2.5 µg/m³\", \"PM10 µg/m³\", \"CO ppm\",\" NO₂ ppm\", \" SO₂ ppm\", \"O₃ ppm\"]\n",
" self.data = []\n",
"\n",
" for _, row in sheet.iterrows():\n",
" name = row[\"Name\"]\n",
" img_path = os.path.join(folder, f\"{name}.jpg\")\n",
" if os.path.exists(img_path):\n",
" measures = row[self.label_columns].tolist()\n",
" self.data.append((name, measures))\n",
"\n",
" print(f\"Loaded {len(self.data)} valid samples\")\n",
" def __len__(self):\n",
" return len(self.data)\n",
"\n",
" def __getitem__(self, idx):\n",
" name, air_measure = self.data[idx]\n",
" img_path = os.path.join(self.folder, f\"{name}.jpg\")\n",
" img = Image.open(img_path).convert(\"RGB\")\n",
" if self.transform:\n",
" img = self.transform(img)\n",
"\n",
" label = self.convertAQI(air_measure)\n",
" return img, label\n",
"\n",
" def convertAQI (self, air_measure):\n",
"\n",
" def truncate(x, n):\n",
" factor = 10 ** n\n",
" return math.trunc(x * factor) / factor\n",
"\n",
" def AQIcategory(aqi):\n",
" if aqi <= 50: return 0\n",
" if aqi <= 100: return 1\n",
" return 2\n",
"\n",
" PM25 = truncate(air_measure[0], 1)\n",
" PM10 = round(air_measure[1])\n",
" CO = truncate(air_measure[2], 1)\n",
" NO2 = round(air_measure[3])\n",
" SO2 = round(air_measure[4])\n",
" O3 = truncate(air_measure[5], 3)\n",
"\n",
" AQI = []\n",
" pollutants = {\"PM25\": PM25, \"PM10\": PM10, \"CO\": CO, \"NO2\": NO2, \"SO2\": SO2, \"O3\": O3}\n",
"\n",
" for key, value in pollutants.items():\n",
" AQIvalue = self.pollutantAQI(key, value)\n",
" AQI.append(AQIvalue)\n",
"\n",
" imageAQI = round(max(AQI))\n",
" return AQIcategory(imageAQI)\n",
"\n",
" def computeAQI (self, cp, bplo, bphi, ilo, ihi):\n",
" return((((ihi - ilo)/(bphi - bplo))*(cp - bplo)) + ilo)\n",
"\n",
" def pollutantAQI(self, pollutant, cp):\n",
" for bplo, bphi, ilo, ihi in BREAKPOINTS[pollutant]:\n",
" if bplo <= cp <= bphi:\n",
" results = self.computeAQI(cp, bplo, bphi, ilo, ihi)\n",
" break\n",
" return results"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Bqja7ztGy7pb"
},
"outputs": [],
"source": [
"#Transforms for Training and Testing Splits\n",
"train_transform = transforms.Compose([\n",
" transforms.RandomHorizontalFlip(p=0.5),\n",
" transforms.Resize((224, 224)),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406],\n",
" [0.229, 0.224, 0.225])\n",
"])\n",
"\n",
"test_transform = transforms.Compose([\n",
" transforms.Resize((224, 224)),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406],\n",
" [0.229, 0.224, 0.225])\n",
"])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "-QtR3wUH0nR1"
},
"outputs": [],
"source": [
"#Full Training and Testing Splits\n",
"train_dataset = AirQualityDataset(IMAGE_DIR, 'train', train_transform)\n",
"train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
"\n",
"test_dataset = AirQualityDataset(IMAGE_DIR, 'test', test_transform)\n",
"test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rBZLL1h1M3mh"
},
"outputs": [],
"source": [
"#Overfitting Training and Testing Splits\n",
"overfit_train_dataset = AirQualityDataset(IMAGE_DIR, 'overfit', train_transform)\n",
"overfit_train_loader = DataLoader(overfit_train_dataset, batch_size=32, shuffle=True)\n",
"\n",
"overfit_test_dataset = AirQualityDataset(IMAGE_DIR, 'overfit', test_transform)\n",
"overfit_test_loader = DataLoader(overfit_test_dataset, batch_size=32, shuffle=False)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yPnlpLSZBtai"
},
"outputs": [],
"source": [
"class VGG19Classifier(nn.Module):\n",
" def __init__(self, num_classes=3, freeze_backbone=True, hidden_dim=512):\n",
" super(VGG19Classifier, self).__init__()\n",
"\n",
" # Load pretrained VGG19\n",
" vgg19 = models.vgg19(pretrained=True)\n",
"\n",
" if freeze_backbone:\n",
" for param in vgg19.features.parameters():\n",
" param.requires_grad = False\n",
"\n",
" self.features = vgg19.features\n",
" self.avgpool = vgg19.avgpool\n",
" self.num_features = 25088\n",
"\n",
" self.classifier = nn.Sequential(OrderedDict([\n",
" ('fc1', nn.Linear(self.num_features, hidden_dim)),\n",
" ('relu', nn.ReLU()),\n",
" ('dropout1', nn.Dropout(0.5)),\n",
" ('fc2', nn.Linear(hidden_dim, num_classes))\n",
" ]))\n",
"\n",
" def forward(self, x):\n",
" x = self.features(x)\n",
" x = self.avgpool(x)\n",
" x = torch.flatten(x, 1)\n",
" x = self.classifier(x)\n",
" return x"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "kVonLOQxvKWA"
},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"from torchvision import models\n",
"from collections import OrderedDict\n",
"\n",
"\n",
"class EfficientNetB0Classifier(nn.Module):\n",
" def __init__(self, num_classes=4, freeze_backbone=True, hidden_dim=512):\n",
" super().__init__()\n",
"\n",
" # Load pretrained EfficientNet-B0\n",
" effnet = models.efficientnet_b0(pretrained=True)\n",
"\n",
" if freeze_backbone:\n",
" for param in effnet.features.parameters():\n",
" param.requires_grad = False\n",
"\n",
" self.features = effnet.features\n",
" self.avgpool = effnet.avgpool\n",
"\n",
" self.num_features = effnet.classifier[1].in_features\n",
"\n",
" # Custom classifier head\n",
" self.classifier = nn.Sequential(OrderedDict([\n",
" (\"fc1\", nn.Linear(self.num_features, hidden_dim)),\n",
" (\"bn1\", nn.BatchNorm1d(hidden_dim)),\n",
" (\"relu\", nn.ReLU(inplace=True)),\n",
" (\"dropout\", nn.Dropout(0.5)),\n",
" (\"fc2\", nn.Linear(hidden_dim, num_classes)),\n",
" ]))\n",
"\n",
"\n",
" def forward(self, x):\n",
" x = self.features(x)\n",
" x = self.avgpool(x)\n",
" x = torch.flatten(x, 1)\n",
" x = self.classifier(x)\n",
" return x\n"
]
},
{
"cell_type": "code",
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
"\n",
"\n",
"class ConvBlock(nn.Module):\n",
" def __init__(self, in_ch, out_ch, stride=1):\n",
" super().__init__()\n",
" self.block = nn.Sequential(\n",
" nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=stride, padding=1, bias=False),\n",
" nn.BatchNorm2d(out_ch),\n",
" nn.ReLU(inplace=True),\n",
" )\n",
"\n",
" def forward(self, x):\n",
" return self.block(x)\n",
"\n",
"\n",
"class AirQualityNet(nn.Module):\n",
" def __init__(self, num_classes=3, dropout=0.5):\n",
" super().__init__()\n",
"\n",
" # -------- Backbone --------\n",
" self.stage1 = nn.Sequential(\n",
" ConvBlock(3, 32),\n",
" ConvBlock(32, 32),\n",
" nn.MaxPool2d(2),\n",
" )\n",
"\n",
" self.stage2 = nn.Sequential(\n",
" ConvBlock(32, 64),\n",
" ConvBlock(64, 64),\n",
" nn.MaxPool2d(2),\n",
" )\n",
"\n",
" self.stage3 = nn.Sequential(\n",
" ConvBlock(64, 128),\n",
" ConvBlock(128, 128),\n",
" nn.MaxPool2d(2),\n",
" )\n",
"\n",
" self.stage4 = nn.Sequential(\n",
" ConvBlock(128, 256),\n",
" nn.MaxPool2d(2),\n",
" )\n",
"\n",
" # -------- Head --------\n",
" self.global_pool = nn.AdaptiveAvgPool2d(1)\n",
"\n",
" self.classifier = nn.Sequential(\n",
" nn.Linear(256, 128),\n",
" nn.ReLU(inplace=True),\n",
" nn.Dropout(dropout),\n",
" nn.Linear(128, num_classes),\n",
" )\n",
"\n",
" def forward(self, x):\n",
" x = self.stage1(x)\n",
" x = self.stage2(x)\n",
" x = self.stage3(x)\n",
" x = self.stage4(x)\n",
"\n",
" x = self.global_pool(x)\n",
" x = torch.flatten(x, 1)\n",
" x = self.classifier(x)\n",
" return x\n"
],
"metadata": {
"id": "3HPR5098uf6h"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"class ResNet101Classifier(nn.Module):\n",
"\n",
" def __init__(self, num_classes=4, freeze_backbone=True, hidden_dim=512):\n",
" super(ResNet101Classifier,self).__init__()\n",
" resnet101 = models.resnet101(pretrained=True)\n",
"\n",
" for param in resnet101.parameters():\n",
" param.requires_grad = False\n",
"\n",
" self.features = nn.Sequential(*list(resnet101.children())[:-1])\n",
" self.num_features = 2048\n",
" self.classifier = nn.Sequential(OrderedDict([ ('fc1',nn.Linear(self.num_features, hidden_dim)), ('relu', nn.ReLU()), ('fc2', nn.Linear(hidden_dim, num_classes)) ]))\n",
"\n",
" def forward(self, x):\n",
" x = self.features(x)\n",
" x = torch.flatten(x, 1)\n",
" x = self.classifier(x)\n",
" return x"
],
"metadata": {
"id": "npaDg3VHuwMa"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "MvmZvEDDvHAz"
},
"outputs": [],
"source": [
"model = EfficientNetB0Classifier().cuda()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HdKCIFYMvxhD"
},
"source": [
"**Loops**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "CcrXoU_nwOQY"
},
"outputs": [],
"source": [
"model = EfficientNetB0Classifier(num_classes=3).cuda()\n",
"criterion = nn.CrossEntropyLoss()\n",
"optimizer = torch.optim.Adam(model.classifier.parameters(), lr=1e-4)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "MK7GEHCbBTsW"
},
"outputs": [],
"source": [
"loss = nn.CrossEntropyLoss()\n",
"input = torch.randn(3, 5, requires_grad=True)\n",
"target = torch.empty(3, dtype=torch.long).random_(5)\n",
"output = loss(input, target)\n",
"output.backward()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "P7OTnRFTwbC-"
},
"outputs": [],
"source": [
"def training_loop(loader, epoch, criterion, model, optimizer, device=\"cuda\"):\n",
" model.train()\n",
"\n",
" total_loss = 0.0\n",
" correct = 0\n",
" total = 0\n",
"\n",
" for step, (images, labels) in enumerate(loader):\n",
" images = images.to(device)\n",
" labels = labels.to(device)\n",
"\n",
" optimizer.zero_grad()\n",
"\n",
" outputs = model(images)\n",
" loss = criterion(outputs, labels)\n",
"\n",
" loss.backward()\n",
" optimizer.step()\n",
"\n",
" batch_size = labels.size(0)\n",
" total_loss += loss.item() * batch_size\n",
" total += batch_size\n",
"\n",
" preds = outputs.argmax(dim=1)\n",
" correct += (preds == labels).sum().item()\n",
"\n",
" if step % 5 == 0:\n",
" print(\n",
" f\"Epoch: {epoch}, Step: {step}, \"\n",
" f\"Loss: {total_loss/total:.4f}, \"\n",
" f\"Accuracy: {(correct/total)*100:.4f}\"\n",
" )\n",
"\n",
" return total_loss / total, correct / total"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5paTe4DDwb4y"
},
"outputs": [],
"source": [
"def validation_loop(loader, epoch, criterion, model, device=\"cuda\"):\n",
" model.eval()\n",
" total_loss = 0\n",
" correct = 0\n",
" total = 0\n",
"\n",
" with torch.no_grad():\n",
" for step, (images, labels) in enumerate(loader):\n",
" images = images.to(device)\n",
" labels = labels.to(device)\n",
"\n",
" outputs = model(images)\n",
" loss = criterion(outputs, labels)\n",
"\n",
" batch_size = labels.size(0)\n",
" total_loss += loss.item() * batch_size\n",
" total += batch_size\n",
"\n",
" preds = outputs.argmax(dim=1)\n",
" correct += (preds == labels).sum().item()\n",
"\n",
" if step % 5 == 0:\n",
" print(\n",
" f\"Epoch: {epoch}, Step: {step}, \"\n",
" f\"Loss: {total_loss/total:.4f}, \"\n",
" f\"Accuracy: {(correct/total)*100:.4f}\"\n",
" )\n",
"\n",
" return total_loss / total, correct / total\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "hDJeTGZ4DJVR"
},
"outputs": [],
"source": [
"epochs = 100\n",
"best_acc = 0.0\n",
"\n",
"checkpoint_path = \"/content/drive/MyDrive/checkpoint.pth\"\n",
"best_model_path = \"/content/drive/MyDrive/best_model.pth\"\n",
"\n",
"start_epoch = 0\n",
"\n",
"train_losses = []\n",
"train_accs = []\n",
"test_losses = []\n",
"test_accs = []\n",
"\n",
"for epoch in range(start_epoch, epochs):\n",
" train_loss, train_acc = training_loop(overfit_train_loader, epoch, criterion, model, optimizer)\n",
" train_losses.append(train_loss)\n",
" train_accs.append(train_acc)\n",
"\n",
" print(f\"Train Epoch: {epoch}, train accuracy: {100*train_acc:.4f}, train loss: {train_loss:.4f}\")\n",
"\n",
" test_loss, test_acc = validation_loop(overfit_test_loader, epoch, criterion, model)\n",
" test_losses.append(test_loss)\n",
" test_accs.append(test_acc)\n",
"\n",
" print(f\"Test Epoch: {epoch}, test accuracy: {100*test_acc:.4f}, test loss: {test_loss:.4f}\")\n",
"\n",
" if test_acc > best_acc:\n",
" best_acc = test_acc\n",
" torch.save(model.state_dict(), best_model_path)\n",
"\n",
" torch.save({\n",
" \"epoch\": epoch,\n",
" \"model\": model.state_dict(),\n",
" \"optimizer\": optimizer.state_dict(),\n",
" \"best_acc\": best_acc\n",
" }, checkpoint_path)\n",
"\n",
"plt.figure(figsize=(12,5))\n",
"plt.subplot(1,2,1)\n",
"plt.plot(train_losses, label='Train Loss')\n",
"plt.plot(test_losses, label='Test Loss')\n",
"plt.xlabel(\"Epoch\")\n",
"plt.ylabel(\"Loss\")\n",
"plt.title(\"Loss over Epochs\")\n",
"plt.legend()\n",
"\n",
"plt.subplot(1,2,2)\n",
"plt.plot([t*100 for t in train_accs], label='Train Accuracy')\n",
"plt.plot([t*100 for t in test_accs], label='Test Accuracy')\n",
"plt.xlabel(\"Epoch\")\n",
"plt.ylabel(\"Accuracy (%)\")\n",
"plt.title(\"Accuracy over Epochs\")\n",
"plt.legend()\n",
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "hKF2PM2wx8NN"
},
"outputs": [],
"source": [
"model.load_state_dict(torch.load(\"/content/drive/MyDrive/best_model.pth\"))\n",
"\n",
"final_loss, final_acc = validation_loop(test_loader, 1, criterion, model)\n",
"print(f\"Final Test Loss: {final_loss:.4f}\")\n",
"print(f\"Final Test Accuracy: {final_acc:.4f}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "gKHNTsCV_8Jo"
},
"outputs": [],
"source": [
"from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay\n",
"\n",
"y_true = []\n",
"y_pred = []\n",
"\n",
"test_loader = DataLoader(overfit_test_dataset, batch_size=1, shuffle=False)\n",
"\n",
"model.eval()\n",
"\n",
"with torch.no_grad():\n",
" for image, label in overfit_test_loader:\n",
" image = image.cuda()\n",
"\n",
" outputs = model(image)\n",
" preds = outputs.argmax(dim=1)\n",
"\n",
" y_true.append(label.cpu().numpy())\n",
" y_pred.append(preds.cpu().numpy())\n",
"\n",
"numpy_true = np.array(y_true)\n",
"numpy_pred = np.array(y_pred)\n",
"\n",
"cm = confusion_matrix(numpy_true, numpy_pred, normalize = \"true\")\n",
"disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels = np.array([0, 1, 2]))\n",
"\n",
"disp.plot()\n",
"plt.show()"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"provenance": [],
"include_colab_link": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment