nn_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NN Training\n",
    "\n",
    "Plan:\n",
    "\n",
    "* Fine-tune a pretrained image classification model\n",
    "    * Compare a few small efficient ones\n",
    "    * Use old PyTorch utility functions\n",
    "    * If enough time, redo in Lightning and write utility functions for it\n",
    "\n",
    "\n",
    "[Nets in torchvision](https://pytorch.org/vision/stable/models.html), though could look at those through timm or directly on huggingface.\n",
    "\n",
    "The image resolution (224x224) matches many standard nets (e.g. [EfficientNet](https://pytorch.org/vision/stable/models/efficientnet.html) B0, [ResNet](https://pytorch.org/vision/stable/models/resnet.html)(s) 34 and 50).\n",
    "The [RexNet family](https://github.com/clovaai/rexnet) is also all of that resolution, and although less common, supposedly more efficient in training than EfficientNet.\n",
    "\n",
    "EfficientNet B0 and RexNet 1.0 have around 5M parameters, and EfficientNet B2 and RexNet 1.5 have around 10M.\n",
    "If venturing farther, then [EfficientNetV2](https://pytorch.org/vision/stable/models/efficientnetv2.html) (the small one, which would still upscale to 384x384) would be even more efficient, but at around 20M parameters (between those of EfficientNet B4 and B5)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import sys\n",
    "# sys.path.append('..')\n",
    "from pytorch_utils import *\n",
    "from lightning_utils import *\n",
    "from pytorch_vision_utils import *\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = r'E:\\Data_and_Models\\Kaggle_Cards'\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pure PyTorch Version\n",
    "\n",
    "### Model Creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The HWC -> CHW permutation seems to happen on its own (or the data info on Kaggle is wrong)\n",
    "# dataloaders 0, 1, 2 are train, test and valid\n",
    "\n",
    "# # EfficientNet B0 - 5.3M parameters\n",
    "# dataloaders, classes = image_dataloaders(data_path, (transforms := (weights := tv.models.EfficientNet_B0_Weights.DEFAULT).transforms()), batch_size = 32)\n",
    "# model = tv.models.efficientnet_b0(weights = weights).to(device)\n",
    "\n",
    "# # EfficientNet B2 - 9.2M parameters\n",
    "# dataloaders, classes = image_dataloaders(data_path, (transforms := (weights := tv.models.EfficientNet_B2_Weights.DEFAULT).transforms()), batch_size = 32)\n",
    "# model = tv.models.efficientnet_b2(weights = weights).to(device)\n",
    "\n",
    "# RexNet 1.0 - 4.8M parameters - https://huggingface.co/timm/rexnet_100.nav_in1k\n",
    "model = timm.create_model('rexnet_100.nav_in1k', pretrained = True, num_classes = 53).eval().to(device) # Cannot use len(classes) yet\n",
    "dataloaders, classes = image_dataloaders(data_path, (transforms := timm.data.create_transform(**timm.data.resolve_model_data_config(model), is_training = False)), batch_size = 32)\n",
    "\n",
    "# # RexNet 1.5 - 9.7M parameters - https://huggingface.co/timm/rexnet_150.nav_in1k\n",
    "# model = timm.create_model('rexnet_150.nav_in1k', pretrained = True, num_classes = 53).eval().to(device) # Cannot use len(classes) yet\n",
    "# dataloaders, classes = image_dataloaders(data_path, (transforms := timm.data.create_transform(**timm.data.resolve_model_data_config(model), is_training = False)), batch_size = 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "=============================================================================================================================\n",
       "Layer (type (var_name))                       Input Shape          Output Shape         Param #              Trainable\n",
       "=============================================================================================================================\n",
       "RexNet (RexNet)                               [32, 3, 224, 224]    [32, 53]             --                   True\n",
       "├─ConvNormAct (stem)                          [32, 3, 224, 224]    [32, 32, 112, 112]   --                   True\n",
       "│    └─Conv2d (conv)                          [32, 3, 224, 224]    [32, 32, 112, 112]   864                  True\n",
       "│    └─BatchNormAct2d (bn)                    [32, 32, 112, 112]   [32, 32, 112, 112]   64                   True\n",
       "│    │    └─Identity (drop)                   [32, 32, 112, 112]   [32, 32, 112, 112]   --                   --\n",
       "│    │    └─SiLU (act)                        [32, 32, 112, 112]   [32, 32, 112, 112]   --                   --\n",
       "├─Sequential (features)                       [32, 32, 112, 112]   [32, 1280, 7, 7]     --                   True\n",
       "│    └─LinearBottleneck (0)                   [32, 32, 112, 112]   [32, 16, 112, 112]   --                   True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 32, 112, 112]   [32, 32, 112, 112]   352                  True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 32, 112, 112]   [32, 32, 112, 112]   --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 32, 112, 112]   [32, 16, 112, 112]   544                  True\n",
       "│    └─LinearBottleneck (1)                   [32, 16, 112, 112]   [32, 27, 56, 56]     --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 16, 112, 112]   [32, 96, 112, 112]   1,728                True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 96, 112, 112]   [32, 96, 56, 56]     1,056                True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 96, 56, 56]     [32, 96, 56, 56]     --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 96, 56, 56]     [32, 27, 56, 56]     2,646                True\n",
       "│    └─LinearBottleneck (2)                   [32, 27, 56, 56]     [32, 38, 56, 56]     --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 27, 56, 56]     [32, 162, 56, 56]    4,698                True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 162, 56, 56]    [32, 162, 56, 56]    1,782                True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 162, 56, 56]    [32, 162, 56, 56]    --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 162, 56, 56]    [32, 38, 56, 56]     6,232                True\n",
       "│    └─LinearBottleneck (3)                   [32, 38, 56, 56]     [32, 50, 28, 28]     --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 38, 56, 56]     [32, 228, 56, 56]    9,120                True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 228, 56, 56]    [32, 228, 28, 28]    2,508                True\n",
       "│    │    └─SEModule (se)                     [32, 228, 28, 28]    [32, 228, 28, 28]    8,949                True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 228, 28, 28]    [32, 228, 28, 28]    --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 228, 28, 28]    [32, 50, 28, 28]     11,500               True\n",
       "│    └─LinearBottleneck (4)                   [32, 50, 28, 28]     [32, 61, 28, 28]     --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 50, 28, 28]     [32, 300, 28, 28]    15,600               True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 300, 28, 28]    [32, 300, 28, 28]    3,300                True\n",
       "│    │    └─SEModule (se)                     [32, 300, 28, 28]    [32, 300, 28, 28]    15,375               True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 300, 28, 28]    [32, 300, 28, 28]    --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 300, 28, 28]    [32, 61, 28, 28]     18,422               True\n",
       "│    └─LinearBottleneck (5)                   [32, 61, 28, 28]     [32, 72, 14, 14]     --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 61, 28, 28]     [32, 366, 28, 28]    23,058               True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 366, 28, 28]    [32, 366, 14, 14]    4,026                True\n",
       "│    │    └─SEModule (se)                     [32, 366, 14, 14]    [32, 366, 14, 14]    22,416               True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 366, 14, 14]    [32, 366, 14, 14]    --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 366, 14, 14]    [32, 72, 14, 14]     26,496               True\n",
       "│    └─LinearBottleneck (6)                   [32, 72, 14, 14]     [32, 84, 14, 14]     --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 72, 14, 14]     [32, 432, 14, 14]    31,968               True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 432, 14, 14]    [32, 432, 14, 14]    4,752                True\n",
       "│    │    └─SEModule (se)                     [32, 432, 14, 14]    [32, 432, 14, 14]    31,644               True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 432, 14, 14]    [32, 432, 14, 14]    --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 432, 14, 14]    [32, 84, 14, 14]     36,456               True\n",
       "│    └─LinearBottleneck (7)                   [32, 84, 14, 14]     [32, 95, 14, 14]     --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 84, 14, 14]     [32, 504, 14, 14]    43,344               True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 504, 14, 14]    [32, 504, 14, 14]    5,544                True\n",
       "│    │    └─SEModule (se)                     [32, 504, 14, 14]    [32, 504, 14, 14]    42,966               True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 504, 14, 14]    [32, 504, 14, 14]    --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 504, 14, 14]    [32, 95, 14, 14]     48,070               True\n",
       "│    └─LinearBottleneck (8)                   [32, 95, 14, 14]     [32, 106, 14, 14]    --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 95, 14, 14]     [32, 570, 14, 14]    55,290               True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 570, 14, 14]    [32, 570, 14, 14]    6,270                True\n",
       "│    │    └─SEModule (se)                     [32, 570, 14, 14]    [32, 570, 14, 14]    54,291               True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 570, 14, 14]    [32, 570, 14, 14]    --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 570, 14, 14]    [32, 106, 14, 14]    60,632               True\n",
       "│    └─LinearBottleneck (9)                   [32, 106, 14, 14]    [32, 117, 14, 14]    --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 106, 14, 14]    [32, 636, 14, 14]    68,688               True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 636, 14, 14]    [32, 636, 14, 14]    6,996                True\n",
       "│    │    └─SEModule (se)                     [32, 636, 14, 14]    [32, 636, 14, 14]    68,211               True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 636, 14, 14]    [32, 636, 14, 14]    --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 636, 14, 14]    [32, 117, 14, 14]    74,646               True\n",
       "│    └─LinearBottleneck (10)                  [32, 117, 14, 14]    [32, 128, 14, 14]    --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 117, 14, 14]    [32, 702, 14, 14]    83,538               True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 702, 14, 14]    [32, 702, 14, 14]    7,722                True\n",
       "│    │    └─SEModule (se)                     [32, 702, 14, 14]    [32, 702, 14, 14]    82,308               True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 702, 14, 14]    [32, 702, 14, 14]    --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 702, 14, 14]    [32, 128, 14, 14]    90,112               True\n",
       "│    └─LinearBottleneck (11)                  [32, 128, 14, 14]    [32, 140, 7, 7]      --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 128, 14, 14]    [32, 768, 14, 14]    99,840               True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 768, 14, 14]    [32, 768, 7, 7]      8,448                True\n",
       "│    │    └─SEModule (se)                     [32, 768, 7, 7]      [32, 768, 7, 7]      99,264               True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 768, 7, 7]      [32, 768, 7, 7]      --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 768, 7, 7]      [32, 140, 7, 7]      107,800              True\n",
       "│    └─LinearBottleneck (12)                  [32, 140, 7, 7]      [32, 151, 7, 7]      --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 140, 7, 7]      [32, 840, 7, 7]      119,280              True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 840, 7, 7]      [32, 840, 7, 7]      9,240                True\n",
       "│    │    └─SEModule (se)                     [32, 840, 7, 7]      [32, 840, 7, 7]      118,650              True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 840, 7, 7]      [32, 840, 7, 7]      --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 840, 7, 7]      [32, 151, 7, 7]      127,142              True\n",
       "│    └─LinearBottleneck (13)                  [32, 151, 7, 7]      [32, 162, 7, 7]      --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 151, 7, 7]      [32, 906, 7, 7]      138,618              True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 906, 7, 7]      [32, 906, 7, 7]      9,966                True\n",
       "│    │    └─SEModule (se)                     [32, 906, 7, 7]      [32, 906, 7, 7]      137,031              True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 906, 7, 7]      [32, 906, 7, 7]      --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 906, 7, 7]      [32, 162, 7, 7]      147,096              True\n",
       "│    └─LinearBottleneck (14)                  [32, 162, 7, 7]      [32, 174, 7, 7]      --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 162, 7, 7]      [32, 972, 7, 7]      159,408              True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 972, 7, 7]      [32, 972, 7, 7]      10,692               True\n",
       "│    │    └─SEModule (se)                     [32, 972, 7, 7]      [32, 972, 7, 7]      158,679              True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 972, 7, 7]      [32, 972, 7, 7]      --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 972, 7, 7]      [32, 174, 7, 7]      169,476              True\n",
       "│    └─LinearBottleneck (15)                  [32, 174, 7, 7]      [32, 185, 7, 7]      --                   True\n",
       "│    │    └─ConvNormAct (conv_exp)            [32, 174, 7, 7]      [32, 1044, 7, 7]     183,744              True\n",
       "│    │    └─ConvNormAct (conv_dw)             [32, 1044, 7, 7]     [32, 1044, 7, 7]     11,484               True\n",
       "│    │    └─SEModule (se)                     [32, 1044, 7, 7]     [32, 1044, 7, 7]     182,961              True\n",
       "│    │    └─ReLU6 (act_dw)                    [32, 1044, 7, 7]     [32, 1044, 7, 7]     --                   --\n",
       "│    │    └─ConvNormAct (conv_pwl)            [32, 1044, 7, 7]     [32, 185, 7, 7]      193,510              True\n",
       "│    └─ConvNormAct (16)                       [32, 185, 7, 7]      [32, 1280, 7, 7]     --                   True\n",
       "│    │    └─Conv2d (conv)                     [32, 185, 7, 7]      [32, 1280, 7, 7]     236,800              True\n",
       "│    │    └─BatchNormAct2d (bn)               [32, 1280, 7, 7]     [32, 1280, 7, 7]     2,560                True\n",
       "├─ClassifierHead (head)                       [32, 1280, 7, 7]     [32, 53]             --                   True\n",
       "│    └─SelectAdaptivePool2d (global_pool)     [32, 1280, 7, 7]     [32, 1280]           --                   --\n",
       "│    │    └─AdaptiveAvgPool2d (pool)          [32, 1280, 7, 7]     [32, 1280, 1, 1]     --                   --\n",
       "│    │    └─Flatten (flatten)                 [32, 1280, 1, 1]     [32, 1280]           --                   --\n",
       "│    └─Dropout (drop)                         [32, 1280]           [32, 1280]           --                   --\n",
       "│    └─Linear (fc)                            [32, 1280]           [32, 53]             67,893               True\n",
       "│    └─Identity (flatten)                     [32, 53]             [32, 53]             --                   --\n",
       "=============================================================================================================================\n",
       "Total params: 3,583,766\n",
       "Trainable params: 3,583,766\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.GIGABYTES): 12.71\n",
       "=============================================================================================================================\n",
       "Input size (MB): 19.27\n",
       "Forward/backward pass size (MB): 1904.75\n",
       "Params size (MB): 14.18\n",
       "Estimated Total Size (MB): 1938.20\n",
       "============================================================================================================================="
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summ(model, input_size = (32, 3, 224, 224))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optionally freeze feature extraction layers\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make the feature extractor layers (\"features\" in the summary) non-trainable (re-run summary above to check)\n",
    "for param in model.features.parameters(): param.requires_grad = False\n",
    "\n",
    "# # This as well for RexNet models\n",
    "for param in model.stem.parameters(): param.requires_grad = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Classifier Layers Replacement (already done for RexNets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Dropout(p=0.2, inplace=True)\n",
       "  (1): Linear(in_features=1280, out_features=1000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# NOT NEEDED FOR REXNETS\n",
    "# Inspect the classifier layers to replicate its structure\n",
    "model.classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOT NEEDED FOR REXNETS\n",
    "\n",
    "# Set the number of classes to the card ones (and reset the other parameters in the classifier layer)\n",
    "model.classifier = torch.nn.Sequential(\n",
    "    # EfficientNet B0\n",
    "    torch.nn.Dropout(p = 0.2, inplace = True),\n",
    "    torch.nn.Linear(in_features = 1280, out_features = len(classes), bias = True)\n",
    "    # # EfficientNet B2\n",
    "    # torch.nn.Dropout(p = 0.3, inplace = True),\n",
    "    # torch.nn.Linear(in_features = 1408, out_features = len(classes), bias = True)\n",
    ").to(device)\n",
    "\n",
    "\n",
    "# The following simpler option does not complain but fails to replace the actual parameter tensor:\n",
    "#   model.classifier[1].out_features = len(classes)\n",
    "# And replacing just the linear layer might be biased by the pretrained dropout one\n",
    "#   model.classifier[1] = torch.nn.Linear(in_features = 1280, out_features = len(classes), bias = True).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model pipeline functions\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "# Define an extra metric beside the loss\n",
    "f1_fn = torchmetrics.F1Score(task = 'multiclass', num_classes = len(classes)).to(device)\n",
    "accuracy_fn = torchmetrics.Accuracy(task = 'multiclass', num_classes = len(classes)).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training RexNet10_0_FullRetrain_Adam001_10_epochs\n",
      "[INFO] Created SummaryWriter, saving to: runs\\2024-04-08\\Cards\\RexNet10\\0_FullRetrain_Adam001_10_epochs...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f931c6d1562a48a0ae83b80c633a789f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1267643d684f47869767b9910d3216c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 1 training batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "332e4c62c12c4d2f91d3c16305e05d1c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 1 testing batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1 | train_loss: 1.7339 | train_metric: 0.5124 | test_loss: 0.5602 | test_metric: 0.8314\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e77b5ff83695458792454196d1048232",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 2 training batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d35602db7324a5e95c1f7b5c6bf71b7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 2 testing batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 2 | train_loss: 0.5716 | train_metric: 0.8264 | test_loss: 0.2966 | test_metric: 0.9147\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f562daa9284422fa78cac7b4cef818b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 3 training batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "843b8c609e7543caa990073ddf8c7f3a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 3 testing batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 3 | train_loss: 0.3233 | train_metric: 0.9004 | test_loss: 0.2243 | test_metric: 0.9198\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d41a0dbaffa645c2ae13fa7b4878a44c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 4 training batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "17073fb1e6fc4867aba35f9ec2502ac3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 4 testing batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 4 | train_loss: 0.2093 | train_metric: 0.9327 | test_loss: 0.1722 | test_metric: 0.9495\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3b98cd2f213545dbaabec92f9a6cebbc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 5 training batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "03162a99655b4e41a683d7deac43253c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 5 testing batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 5 | train_loss: 0.1696 | train_metric: 0.9484 | test_loss: 0.2843 | test_metric: 0.9217\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c93f0f5d13445a7830f8f0423ef850a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 6 training batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "304c9972ea7d4392be98bac0b1b8153d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 6 testing batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 6 | train_loss: 0.1520 | train_metric: 0.9536 | test_loss: 0.2251 | test_metric: 0.9267\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "03db2568283d4080b53e00bfec2b3cc5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 7 training batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a97b9d3bc6d4d4abdd90179473b0323",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 7 testing batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 7 | train_loss: 0.0993 | train_metric: 0.9685 | test_loss: 0.1524 | test_metric: 0.9668\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a7d99d392b9e4618872a8e344e482cb0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 8 training batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "64225debca2e4b26ab774252eb2e77f1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 8 testing batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 8 | train_loss: 0.0992 | train_metric: 0.9704 | test_loss: 0.1931 | test_metric: 0.9475\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "56d02410191b420c87ec9a78435f3bad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 9 training batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "900413ebc8364141aaefdc698f0b5e5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 9 testing batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 9 | train_loss: 0.0868 | train_metric: 0.9748 | test_loss: 0.2121 | test_metric: 0.9425\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3bd7b597b57248e48fa457597c8b221c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 10 training batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "86d878c25e2c4367b1a3e612d79b8c77",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Epoch 10 testing batches: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 10 | train_loss: 0.0903 | train_metric: 0.9732 | test_loss: 0.2218 | test_metric: 0.9390\n"
     ]
    }
   ],
   "source": [
    "# Train the model\n",
    "\n",
    "model_name = 'RexNet10'\n",
    "extra = '0_FullRetrain_Adam001_10_epochs' # Mimicking the train_combination function format\n",
    "print(f'Training {model_name}_{extra}')\n",
    "\n",
    "set_seeds(42)\n",
    "results = fit(model, train_dataloader = dataloaders[0], test_dataloader = dataloaders[1],\n",
    "    optimiser = torch.optim.Adam(model.parameters(), lr = 0.001), loss_fn = loss_fn,\n",
    "    metric_name_and_fn = ('F1', f1_fn),\n",
    "    # metric_name_and_fn = ('Accuracy', accuracy_fn),\n",
    "    epochs = 10,\n",
    "    writer = tensorboard_writer(experiment_name = 'Cards', model_name = model_name, extra = extra, save_dir = fr'{data_path}\\runs')\n",
    "    # writer = None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] Saving model to: E:\\Data_and_Models\\Kaggle_Cards\\Models\\RexNet10_0_FullRetrain_Adam001_10_epochs.pth\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "WindowsPath('E:/Data_and_Models/Kaggle_Cards/Models/RexNet10_0_FullRetrain_Adam001_10_epochs.pth')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# save_model(model, fr'{data_path}\\Models', f'{model_name}_{extra}.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch Lightning Version\n",
    "\n",
    "### Model Creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = timm.create_model('rexnet_100.nav_in1k', pretrained = True, num_classes = 53)\n",
    "# for param in model.features.parameters(): param.requires_grad = False\n",
    "# for param in model.stem.parameters(): param.requires_grad = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "transforms = timm.data.create_transform(**timm.data.resolve_model_data_config(model), is_training = False)\n",
    "ldata = LocalImageDataModule(data_path, transform = transforms, batch_size = 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model pipeline functions\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "# Define an extra metric beside the loss\n",
    "f1_fn = torchmetrics.F1Score(task = 'multiclass', num_classes = 53)\n",
    "accuracy_fn = torchmetrics.Accuracy(task = 'multiclass', num_classes = 53)\n",
    "\n",
    "# # General prediction pipeline\n",
    "# with torch.inference_mode(): pred_logit = model(transforms(img).unsqueeze(dim = 0).to(device)) # Prepend \"batch\" dimension (-> [batch_size, color_channels, height, width])\n",
    "# pred_prob = torch.softmax(pred_logit, dim = 1)\n",
    "# return torch.argmax(pred_prob, dim = 1)\n",
    "\n",
    "def prediction_fn(logits): return torch.argmax(torch.softmax(logits, dim = 1), dim = 1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==================================================================================================================================\n",
       "Layer (type (var_name))                            Input Shape          Output Shape         Param #              Trainable\n",
       "==================================================================================================================================\n",
       "Strike (Strike)                                    [32, 3, 224, 224]    [32, 53]             --                   Partial\n",
       "├─RexNet (model)                                   [32, 3, 224, 224]    [32, 53]             --                   Partial\n",
       "│    └─ConvNormAct (stem)                          [32, 3, 224, 224]    [32, 32, 112, 112]   --                   False\n",
       "│    │    └─Conv2d (conv)                          [32, 3, 224, 224]    [32, 32, 112, 112]   (864)                False\n",
       "│    │    └─BatchNormAct2d (bn)                    [32, 32, 112, 112]   [32, 32, 112, 112]   (64)                 False\n",
       "│    └─Sequential (features)                       [32, 32, 112, 112]   [32, 1280, 7, 7]     --                   False\n",
       "│    │    └─LinearBottleneck (0)                   [32, 32, 112, 112]   [32, 16, 112, 112]   (896)                False\n",
       "│    │    └─LinearBottleneck (1)                   [32, 16, 112, 112]   [32, 27, 56, 56]     (5,430)              False\n",
       "│    │    └─LinearBottleneck (2)                   [32, 27, 56, 56]     [32, 38, 56, 56]     (12,712)             False\n",
       "│    │    └─LinearBottleneck (3)                   [32, 38, 56, 56]     [32, 50, 28, 28]     (32,077)             False\n",
       "│    │    └─LinearBottleneck (4)                   [32, 50, 28, 28]     [32, 61, 28, 28]     (52,697)             False\n",
       "│    │    └─LinearBottleneck (5)                   [32, 61, 28, 28]     [32, 72, 14, 14]     (75,996)             False\n",
       "│    │    └─LinearBottleneck (6)                   [32, 72, 14, 14]     [32, 84, 14, 14]     (104,820)            False\n",
       "│    │    └─LinearBottleneck (7)                   [32, 84, 14, 14]     [32, 95, 14, 14]     (139,924)            False\n",
       "│    │    └─LinearBottleneck (8)                   [32, 95, 14, 14]     [32, 106, 14, 14]    (176,483)            False\n",
       "│    │    └─LinearBottleneck (9)                   [32, 106, 14, 14]    [32, 117, 14, 14]    (218,541)            False\n",
       "│    │    └─LinearBottleneck (10)                  [32, 117, 14, 14]    [32, 128, 14, 14]    (263,680)            False\n",
       "│    │    └─LinearBottleneck (11)                  [32, 128, 14, 14]    [32, 140, 7, 7]      (315,352)            False\n",
       "│    │    └─LinearBottleneck (12)                  [32, 140, 7, 7]      [32, 151, 7, 7]      (374,312)            False\n",
       "│    │    └─LinearBottleneck (13)                  [32, 151, 7, 7]      [32, 162, 7, 7]      (432,711)            False\n",
       "│    │    └─LinearBottleneck (14)                  [32, 162, 7, 7]      [32, 174, 7, 7]      (498,255)            False\n",
       "│    │    └─LinearBottleneck (15)                  [32, 174, 7, 7]      [32, 185, 7, 7]      (571,699)            False\n",
       "│    │    └─ConvNormAct (16)                       [32, 185, 7, 7]      [32, 1280, 7, 7]     (239,360)            False\n",
       "│    └─ClassifierHead (head)                       [32, 1280, 7, 7]     [32, 53]             --                   True\n",
       "│    │    └─SelectAdaptivePool2d (global_pool)     [32, 1280, 7, 7]     [32, 1280]           --                   --\n",
       "│    │    └─Dropout (drop)                         [32, 1280]           [32, 1280]           --                   --\n",
       "│    │    └─Linear (fc)                            [32, 1280]           [32, 53]             67,893               True\n",
       "│    │    └─Identity (flatten)                     [32, 53]             [32, 53]             --                   --\n",
       "==================================================================================================================================\n",
       "Total params: 3,583,766\n",
       "Trainable params: 67,893\n",
       "Non-trainable params: 3,515,873\n",
       "Total mult-adds (Units.GIGABYTES): 12.71\n",
       "==================================================================================================================================\n",
       "Input size (MB): 19.27\n",
       "Forward/backward pass size (MB): 1904.75\n",
       "Params size (MB): 14.18\n",
       "Estimated Total Size (MB): 1938.20\n",
       "=================================================================================================================================="
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lmod = Strike(model, loss_fn = loss_fn, metric_name_and_fn = ('F1', f1_fn),\n",
    "               optimiser_factory = lambda m: torch.optim.Adam(m.parameters(), lr = m.learning_rate),\n",
    "               prediction_fn = prediction_fn, learning_rate = 0.001, log_at_every_step = False)\n",
    "\n",
    "summ(lmod, input_size = (32, 3, 224, 224))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training FullRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    }
   ],
   "source": [
    "experiment_name = 'FullRetrain_EarlyStop'\n",
    "# experiment_name = 'ClassRetrain_EarlyStop'\n",
    "model_name = 'RexNet10'\n",
    "extra = 'Adam001_max10_epochs' # Mimicking the train_combination function format\n",
    "# extra = 'ClassRetrain_AdamAutoLR_max10_epochs' # Mimicking the train_combination function format\n",
    "print(f'Training {experiment_name}_{model_name}_{extra}')\n",
    "\n",
    "trainer = L.Trainer(\n",
    "    accelerator = 'gpu', devices = 1, \n",
    "    min_epochs = 1, max_epochs = 10,\n",
    "    callbacks = [\n",
    "        EarlyStopping(monitor = 'val_loss', mode = 'min', min_delta = 0.01, patience = 3, verbose = True),\n",
    "        ModelCheckpoint(monitor = 'val_loss', mode = 'min', save_top_k = 3, verbose = True,\n",
    "                        dirpath = (checkpoints_path := fr'{data_path}\\Models\\{experiment_name}_{model_name}_{extra}'),\n",
    "                        filename = '{epoch}-{val_loss:.2f}-{val_F1:.2f}', enable_version_counter = True),\n",
    "                        \n",
    "        ## Unfortunately need to forego automatic lr if want to use early stopping properly (i.e. on validation loss) since the steps are counted and throw off the checkpointing: https://github.com/Lightning-AI/pytorch-lightning/issues/19575\n",
    "        # EarlyStopping(monitor = 'train_F1', min_delta = 0.01, mode = 'max', patience = 3, verbose = True),\n",
    "        # LearningRateFinder(num_training_steps = 100),\n",
    "        # # IteratedLearningRateFinder(at_epochs = [3, 6, 9], num_training_steps = 100),\n",
    "        # # LearningRateMonitor(logging_interval = 'epoch'), # No need if not readjusting it after epoch 0\n",
    "    ],\n",
    "    logger = TBLogger(experiment_name = experiment_name, model_name = model_name, extra = extra, save_dir = fr'{data_path}\\runs'), # The default logger is also good\n",
    "    profiler = PyTorchProfiler(filename = f'{model_name}_{extra}'),\n",
    "\n",
    "    ## Pre-training checks\n",
    "    # fast_dev_run = False # Tries a couple of batches for training, validating and testing first, just to check that everything works\n",
    "    # overfit_batches = 10, # Check able to overfit few batches; int for count or float for proportion\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name      | Type              | Params\n",
      "------------------------------------------------\n",
      "0 | model     | RexNet            | 3.6 M \n",
      "1 | loss_fn   | CrossEntropyLoss  | 0     \n",
      "2 | metric_fn | MulticlassF1Score | 0     \n",
      "------------------------------------------------\n",
      "67.9 K    Trainable params\n",
      "3.5 M     Non-trainable params\n",
      "3.6 M     Total params\n",
      "14.335    Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0f25e8f12ab54af2a3fadebb398e3695",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d3710bf3c83b4bb9bbb506471689bfc0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fa8eb3eafb754cbdb97d023fef452b88",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_loss improved. New best score: 2.933\n",
      "Epoch 0, global step 239: 'val_loss' reached 2.93327 (best 2.93327), saving model to 'E:\\\\Data_and_Models\\\\Kaggle_Cards\\\\Models\\\\ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\\\\epoch=0-val_loss=2.93-val_F1=0.11.ckpt' as top 3\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "60a4a0e6795f42eb94e4ba184cc7d8a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1, global step 478: 'val_loss' reached 3.07336 (best 2.93327), saving model to 'E:\\\\Data_and_Models\\\\Kaggle_Cards\\\\Models\\\\ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\\\\epoch=1-val_loss=3.07-val_F1=0.22.ckpt' as top 3\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "905465ec2baa46d884732a47991274a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2, global step 717: 'val_loss' reached 3.30281 (best 2.93327), saving model to 'E:\\\\Data_and_Models\\\\Kaggle_Cards\\\\Models\\\\ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\\\\epoch=2-val_loss=3.30-val_F1=0.11.ckpt' as top 3\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "386fd848270e450982dd7b5c6fcdf936",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_loss improved by 0.104 >= min_delta = 0.01. New best score: 2.829\n",
      "Epoch 3, global step 956: 'val_loss' reached 2.82916 (best 2.82916), saving model to 'E:\\\\Data_and_Models\\\\Kaggle_Cards\\\\Models\\\\ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\\\\epoch=3-val_loss=2.83-val_F1=0.22.ckpt' as top 3\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a6fb0ccb7ebe404899cb72ca3afc7080",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 4, global step 1195: 'val_loss' reached 2.86953 (best 2.82916), saving model to 'E:\\\\Data_and_Models\\\\Kaggle_Cards\\\\Models\\\\ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\\\\epoch=4-val_loss=2.87-val_F1=0.22.ckpt' as top 3\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b7c3edf41db43feb16d7edbf079a24d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 5, global step 1434: 'val_loss' reached 2.90167 (best 2.82916), saving model to 'E:\\\\Data_and_Models\\\\Kaggle_Cards\\\\Models\\\\ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\\\\epoch=5-val_loss=2.90-val_F1=0.33.ckpt' as top 3\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90515ace7485438bb4ff114a318b87cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_loss improved by 0.016 >= min_delta = 0.01. New best score: 2.813\n",
      "Epoch 6, global step 1673: 'val_loss' reached 2.81331 (best 2.81331), saving model to 'E:\\\\Data_and_Models\\\\Kaggle_Cards\\\\Models\\\\ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\\\\epoch=6-val_loss=2.81-val_F1=0.33.ckpt' as top 3\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc41fbe1b1224d769c3c688a5b2d0f63",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Metric val_loss improved by 0.279 >= min_delta = 0.01. New best score: 2.534\n",
      "Epoch 7, global step 1912: 'val_loss' reached 2.53411 (best 2.53411), saving model to 'E:\\\\Data_and_Models\\\\Kaggle_Cards\\\\Models\\\\ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\\\\epoch=7-val_loss=2.53-val_F1=0.33.ckpt' as top 3\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5e72797603a24c389e9b365b65d5ddea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 8, global step 2151: 'val_loss' reached 2.73769 (best 2.53411), saving model to 'E:\\\\Data_and_Models\\\\Kaggle_Cards\\\\Models\\\\ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs\\\\epoch=8-val_loss=2.74-val_F1=0.22.ckpt' as top 3\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "97cd555278ae489fa2902a109506d878",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 9, global step 2390: 'val_loss' was not in top 3\n",
      "`Trainer.fit` stopped: `max_epochs=10` reached.\n"
     ]
    }
   ],
   "source": [
    "print(f'Training {experiment_name}_{model_name}_{extra}')\n",
    "\n",
    "set_seeds(42)\n",
    "# trainer.fit(lmod, dataloaders[0], dataloaders[2])\n",
    "trainer.fit(lmod, ldata,\n",
    "            # ckpt_path = trainer.checkpoint_callback.best_model_path # Continue training from a checkpoint\n",
    ")\n",
    "\n",
    "# NOTE: trainer will prevent from re-fitting, so if the model changed need to re-declare trainer above in order to train again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "59b55489d24a42caa1677c0355af26e0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "     Validate metric           DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "         val_F1             0.3333333432674408\n",
      "        val_loss            2.5341060161590576\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f945f7c0ddc442d5aac358cf31af081b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "         test_F1            0.2222222238779068\n",
      "        test_loss            2.985792636871338\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 2.985792636871338, 'test_F1': 0.2222222238779068}]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Need to give arguments again since many are non-pickleable with .save_hyperparameters\n",
    "#   Could solve by moving those ones to a function producing the class\n",
    "bestmod = Strike.load_from_checkpoint(trainer.checkpoint_callback.best_model_path,\n",
    "    model = model, loss_fn = loss_fn, metric_name_and_fn = ('F1', f1_fn),\n",
    "    optimiser_factory = lambda m: torch.optim.Adam(m.parameters(), lr = m.learning_rate),\n",
    "    prediction_fn = prediction_fn, learning_rate = 0.001, log_at_every_step = False)\n",
    "\n",
    "trainer.validate(bestmod, ldata)\n",
    "trainer.test(bestmod, ldata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] Saving model to: E:\\Data_and_Models\\Kaggle_Cards\\Models\\ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs.pth\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "WindowsPath('E:/Data_and_Models/Kaggle_Cards/Models/ClassRetrain_EarlyStop_RexNet10_Adam001_max10_epochs.pth')"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "# save_model(bestmod, fr'{data_path}\\Models', f'{experiment_name}_{model_name}_{extra}.pth')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ML11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}