train_colab.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "# mount google drive\n",
    "drive.mount('/content/drive/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install jsonargparse==4.28.0 lightning==2.2.3 seaborn==0.13.2 tabulate==0.9.0 termcolor==2.4.0 torch==2.2.2 torchmetrics==1.3.2 torchvision==0.17.2 wandb==0.16.6 wget==3.2 tqdm==4.66.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cd /content/drive/MyDrive/hyu/aue8088-pa1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PyTorch & Pytorch Lightning\n",
    "from lightning.pytorch.loggers.wandb import WandbLogger\n",
    "from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint\n",
    "from lightning import Trainer\n",
    "import torch\n",
    "\n",
    "# Custom packages\n",
    "from src.dataset import TinyImageNetDatasetModule\n",
    "from src.network import SimpleClassifier\n",
    "import src.config as cfg\n",
    "\n",
    "torch.set_float32_matmul_precision('medium')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SimpleClassifier(\n",
    "    model_name = cfg.MODEL_NAME,\n",
    "    num_classes = cfg.NUM_CLASSES,\n",
    "    optimizer_params = cfg.OPTIMIZER_PARAMS,\n",
    "    scheduler_params = cfg.SCHEDULER_PARAMS,\n",
    ")\n",
    "\n",
    "datamodule = TinyImageNetDatasetModule(\n",
    "    batch_size = cfg.BATCH_SIZE,\n",
    ")\n",
    "\n",
    "wandb_logger = WandbLogger(\n",
    "    project = cfg.WANDB_PROJECT,\n",
    "    save_dir = cfg.WANDB_SAVE_DIR,\n",
    "    entity = cfg.WANDB_ENTITY,\n",
    "    name = cfg.WANDB_NAME,\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    accelerator = cfg.ACCELERATOR,\n",
    "    devices = cfg.DEVICES,\n",
    "    precision = cfg.PRECISION_STR,\n",
    "    max_epochs = cfg.NUM_EPOCHS,\n",
    "    check_val_every_n_epoch = cfg.VAL_EVERY_N_EPOCH,\n",
    "    logger = wandb_logger,\n",
    "    callbacks = [\n",
    "        LearningRateMonitor(logging_interval='epoch'),\n",
    "        ModelCheckpoint(save_top_k=1, monitor='accuracy/val', mode='max'),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Takes time at first (download dataset)\n",
    "trainer.fit(model, datamodule=datamodule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.validate(ckpt_path='best', datamodule=datamodule)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}