quickstart.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0c461fb7-78b9-4201-847e-2e11c8f599ff",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Quick start guide\n",
    "This notebook serves as an example of how to train a simple model using pytorch and the ready-to-train AI4Arctic challenge dataset. Initially, a dictionary, 'train_options', is set up with relevant options for both the example U-Net Convolutional Neural Network model and the dataloader. Note that the weights of the U-Net will be initialised at random and therefore not deterministic - results will vary for every training run. Two lists (dataset.json and testset.json) include the names of the scenes relevant to training and testing, where the former can be altered if desired. Training data is loaded in parallel using the build-in torch Dataset and Dataloader classes, and works by randomly sampling a scene and performing a random crop to extract a patch. Each batch will then be compiled of X number of these patches with the patch size in the 'train_options'. An obstacle is different grid resolution sizes, which is overcome by upsampling low resolution variables, e.g. AMSR2, ERA5, to match the SAR pixels. A number of batches will be prepared in parallel and stored until use, depending on the number of workers (processes) spawned (this can be changed in 'num_workers' in 'train_options'). The model is trained on a fixed number of steps according to the number of batches in an epoch, defined by the 'epoch_len' parameter, and will run for a total number of epochs depending on the 'epochs' parameter. After each epoch, the model is evaluated. In this example, a random number of scenes are sampled among the training scenes (and removed from the list of training scenes) to act as a validation set used for the evaluation. The model is evaluated with the metrics, and if the current validation attempt is superior to the previous, then the model parameters are stored in the 'best_model' file in the directory.\n",
    "\n",
    "The models are scored on the three sea ice parameters; Sea Ice Concentration (SIC), Stage of Development (SOD) and the Floe size (FLOE) with the $R²$ metric for the SIC, and the weighted F1 metric for the SOD and FLOE. The 3 scores are combined into a single metric by taking the weighted average with SIC and SOD being weighted with 2 and the FLOE with 1.\n",
    "\n",
    "Finally, once you are ready to test your model on the test scenes (without reference data), the 'test_upload' notebook will produce model outputs with your model of choice and save the output as a netCDF file, which can be uploaded to the AI4EO.eu website. The model outputs will be evaluated and then you will receive a score. \n",
    "\n",
    "This quick start notebook is by no means necessary to utilize, and you are more than welcome to develop your own data pipeline. We do however require that the model output is stored in a netcdf file with xarray.dataarrays titled '{scene_name}_{chart}', i.e. 3 charts per scene / file (see how in 'test_upload'). In addition, you are more than welcome to create your own preprocessing scheme to prepare the raw AI4Arctic challenge dataset. However, we ask that the model output is in 80 m pixel spacing (original is 40 m), and that you follow the class numberings from the lookup tables in 'utils' - at least you will be evaluated in this way. Furthermore, we have included a function to convert the polygon_icechart to SIC, SOD and FLOE, you will have to incorporate it yourself.\n",
    "\n",
    "The first cell imports the necessary Python packages, initializes the 'train_options' dictionary, the sample U-Net options, loads the dataset list and select validation scenes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "44a8d5a8-674a-42b3-9789-2a04818fa860",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -- Built-in modules -- #\n",
    "import gc\n",
    "import os\n",
    "import sys\n",
    "\n",
    "# -- Environmental variables -- #\n",
    "os.environ['AI4ARCTIC_DATA'] = ''  # Fill in directory for data location.\n",
    "os.environ['AI4ARCTIC_ENV'] = ''  # Fill in directory for environment with Ai4Arctic get-started package. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c82d24f0-233b-41f9-95ef-af0cc0895800",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stored 'train_options' (dict)\n",
      "Options initialised\n"
     ]
    }
   ],
   "source": [
    "# -- Third-part modules -- #\n",
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import xarray as xr\n",
    "from tqdm.notebook import tqdm  # Progress bar\n",
    "\n",
    "# --Proprietary modules -- #\n",
    "from functions import chart_cbar, r2_metric, f1_metric, compute_metrics  # Functions to calculate metrics and show the relevant chart colorbar.\n",
    "from loaders import AI4ArcticChallengeDataset, AI4ArcticChallengeTestDataset, get_variable_options  # Custom dataloaders for regular training and validation.\n",
    "from unet import UNet  # Convolutional Neural Network model\n",
    "from utils import CHARTS, SIC_LOOKUP, SOD_LOOKUP, FLOE_LOOKUP, SCENE_VARIABLES, colour_str\n",
    "\n",
    "train_options = {\n",
    "    # -- Training options -- #\n",
    "    'path_to_processed_data': os.environ['AI4ARCTIC_DATA'],  # Replace with data directory path.\n",
    "    'path_to_env': os.environ['AI4ARCTIC_ENV'],  # Replace with environmment directory path.\n",
    "    'lr': 0.0001,  # Optimizer learning rate.\n",
    "    'epochs': 50,  # Number of epochs before training stop.\n",
    "    'epoch_len': 500,  # Number of batches for each epoch.\n",
    "    'patch_size': 256,  # Size of patches sampled. Used for both Width and Height.\n",
    "    'batch_size': 8,  # Number of patches for each batch.\n",
    "    'loader_upsampling': 'nearest',  # How to upscale low resolution variables to high resolution.\n",
    "    \n",
    "    # -- Data prepraration lookups and metrics.\n",
    "    'train_variables': SCENE_VARIABLES,  # Contains the relevant variables in the scenes.\n",
    "    'charts': CHARTS,  # Charts to train on.\n",
    "    'n_classes': {  # number of total classes in the reference charts, including the mask.\n",
    "        'SIC': SIC_LOOKUP['n_classes'],\n",
    "        'SOD': SOD_LOOKUP['n_classes'],\n",
    "        'FLOE': FLOE_LOOKUP['n_classes']\n",
    "    },\n",
    "    'pixel_spacing': 80,  # SAR pixel spacing. 80 for the ready-to-train AI4Arctic Challenge dataset.\n",
    "    'train_fill_value': 0,  # Mask value for SAR training data.\n",
    "    'class_fill_values': {  # Mask value for class/reference data.\n",
    "        'SIC': SIC_LOOKUP['mask'],\n",
    "        'SOD': SOD_LOOKUP['mask'],\n",
    "        'FLOE': FLOE_LOOKUP['mask'],\n",
    "    },\n",
    "    \n",
    "    # -- Validation options -- #\n",
    "    'chart_metric': {  # Metric functions for each ice parameter and the associated weight.\n",
    "        'SIC': {\n",
    "            'func': r2_metric,\n",
    "            'weight': 2,\n",
    "        },\n",
    "        'SOD': {\n",
    "            'func': f1_metric,\n",
    "            'weight': 2,\n",
    "        },\n",
    "        'FLOE': {\n",
    "            'func': f1_metric,\n",
    "            'weight': 1,\n",
    "        },\n",
    "    },\n",
    "    'num_val_scenes': 10,  # Number of scenes randomly sampled from train_list to use in validation.\n",
    "    \n",
    "    # -- GPU/cuda options -- #\n",
    "    'gpu_id': 0,  # Index of GPU. In case of multiple GPUs.\n",
    "    'num_workers': 6,  # Number of parallel processes to fetch data.\n",
    "    'num_workers_val': 1,  # Number of parallel processes during validation.\n",
    "    \n",
    "    # -- U-Net Options -- #\n",
    "    'unet_conv_filters': [16, 32, 64, 64],  # Number of filters in the U-Net.\n",
    "    'conv_kernel_size': (3, 3),  # Size of convolutional kernels.\n",
    "    'conv_stride_rate': (1, 1),  # Stride rate of convolutional kernels.\n",
    "    'conv_dilation_rate': (1, 1),  # Dilation rate of convolutional kernels.\n",
    "    'conv_padding': (1, 1),  # Number of padded pixels in convolutional layers.\n",
    "    'conv_padding_style': 'zeros',  # Style of padding.\n",
    "}\n",
    "# Get options for variables, amsrenv grid, cropping and upsampling.\n",
    "get_variable_options = get_variable_options(train_options)\n",
    "# To be used in test_upload.\n",
    "%store train_options  \n",
    "\n",
    "# Load training list.\n",
    "with open(train_options['path_to_env'] + 'datalists/dataset.json') as file:\n",
    "    train_options['train_list'] = json.loads(file.read())\n",
    "# Convert the original scene names to the preprocessed names.\n",
    "train_options['train_list'] = [file[17:32] + '_' + file[77:80] + '_prep.nc' for file in train_options['train_list']]\n",
    "# Select a random number of validation scenes with the same seed. Feel free to change the seed.et\n",
    "np.random.seed(0)\n",
    "train_options['validate_list'] = np.random.choice(np.array(train_options['train_list']), size=train_options['num_val_scenes'], replace=False)\n",
    "# Remove the validation scenes from the train list.\n",
    "train_options['train_list'] = [scene for scene in train_options['train_list'] if scene not in train_options['validate_list']]\n",
    "print('Options initialised')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74517e22-5636-4004-84c5-3cc416276054",
   "metadata": {
    "tags": []
   },
   "source": [
    "### CUDA / GPU Setup\n",
    "This sets up the 'device' variable containing GPU information, and the custom dataset and dataloader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "26c3c79a-3f60-4ca3-a6e9-b967929a3c12",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0;32mGPU available!\u001b[0m\n",
      "Total number of available devices:  \u001b[0;33m4\u001b[0m\n",
      "GPU and data setup complete.\n"
     ]
    }
   ],
   "source": [
    "# Get GPU resources.\n",
    "if torch.cuda.is_available():\n",
    "    print(colour_str('GPU available!', 'green'))\n",
    "    print('Total number of available devices: ', colour_str(torch.cuda.device_count(), 'orange'))\n",
    "    device = torch.device(f\"cuda:{train_options['gpu_id']}\")\n",
    "\n",
    "else:\n",
    "    print(colour_str('GPU not available.', 'red'))\n",
    "    device = torch.device('cpu')\n",
    "\n",
    "# Custom dataset and dataloader.\n",
    "dataset = AI4ArcticChallengeDataset(files=train_options['train_list'], options=train_options)\n",
    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=None, shuffle=True, num_workers=train_options['num_workers'], pin_memory=True)\n",
    "# - Setup of the validation dataset/dataloader. The same is used for model testing in 'test_upload.ipynb'.\n",
    "dataset_val = AI4ArcticChallengeTestDataset(options=train_options, files=train_options['validate_list'])\n",
    "dataloader_val = torch.utils.data.DataLoader(dataset_val, batch_size=None, num_workers=train_options['num_workers_val'], shuffle=False)\n",
    "\n",
    "print('GPU and data setup complete.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31e9849d-da01-402e-a79f-2f196618fb21",
   "metadata": {},
   "source": [
    "### Example of Model, optimiser and loss function setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "df663da8-9779-4f8e-b641-29c9ea4e6036",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model setup complete\n"
     ]
    }
   ],
   "source": [
    "# Setup U-Net model, adam optimizer, loss function and dataloader.\n",
    "net = UNet(options=train_options).to(device)\n",
    "optimizer = torch.optim.Adam(list(net.parameters()), lr=train_options['lr'])\n",
    "torch.backends.cudnn.benchmark = True  # Selects the kernel with the best performance for the GPU and given input size.\n",
    "\n",
    "# Loss functions to use for each sea ice parameter.\n",
    "# The ignore_index argument discounts the masked values, ensuring that the model is not using these pixels to train on.\n",
    "# It is equivalent to multiplying the loss of the relevant masked pixel with 0.\n",
    "loss_functions = {chart: torch.nn.CrossEntropyLoss(ignore_index=train_options['class_fill_values'][chart]) \\\n",
    "                                                   for chart in train_options['charts']}\n",
    "print('Model setup complete')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc2980de-1025-4b75-9c43-0c3023a2d12c",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Example of model training and validation loop\n",
    "A simple model training loop following by a simple validation loop. Validation is carried out on full scenes, i.e. no cropping or stitching. If there is not enough space on the GPU, then try to do it on the cpu. This can be done by using 'net = net.cpu()'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0d00e418-ff77-4948-92bd-0b2313b281d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.01687455177307129,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 50,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "117514a7ebab417ab371f356f9660f9a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.01271677017211914,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b01fab75e1bf44ec999c7728394e88ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 5.106\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.014016866683959961,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2b2aa7b6fe81450daae22a898abef6cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 3.861\n",
      "Epoch 0 score:\n",
      "SIC r2_metric: 63.423%\n",
      "SOD f1_metric: 82.912%\n",
      "FLOE f1_metric: 66.816%\n",
      "Combined score: 71.897%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.012796878814697266,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0af83ddfe491468ea806bc7c8cc5e515",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 3.705\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.015100955963134766,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "51dbcebf55e24d36b4ffb0de472d82b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 2.776\n",
      "Epoch 1 score:\n",
      "SIC r2_metric: 69.486%\n",
      "SOD f1_metric: 83.712%\n",
      "FLOE f1_metric: 66.73%\n",
      "Combined score: 74.625%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.013161420822143555,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "440a6d2376134917b0663b320239e1ae",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 2.938\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.013636112213134766,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1dda273df30c48f0903dcce440406f9f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 2.968\n",
      "Epoch 2 score:\n",
      "SIC r2_metric: 56.93%\n",
      "SOD f1_metric: 85.064%\n",
      "FLOE f1_metric: 67.399%\n",
      "Combined score: 70.277%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.01253366470336914,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f2f16704771c42c083efd41c1d357df0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 2.517\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.015273571014404297,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ebdd31b02a4a44b5866fc298b4b010bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 0.898\n",
      "Epoch 3 score:\n",
      "SIC r2_metric: 62.42%\n",
      "SOD f1_metric: 81.326%\n",
      "FLOE f1_metric: 67.61%\n",
      "Combined score: 71.02%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.012943267822265625,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1caf1da118104507a90ec0d09e6f80ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 2.269\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.01524662971496582,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b2a21c60da1146f28a17d094b92cad1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 2.526\n",
      "Epoch 4 score:\n",
      "SIC r2_metric: 61.844%\n",
      "SOD f1_metric: 83.445%\n",
      "FLOE f1_metric: 67.758%\n",
      "Combined score: 71.667%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.012260913848876953,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "10ef5c86be0145f899cd26d911c7718d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 2.157\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.016118526458740234,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9ba305283b5b4df2adbf5451ed9c3686",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 2.634\n",
      "Epoch 5 score:\n",
      "SIC r2_metric: 60.616%\n",
      "SOD f1_metric: 85.164%\n",
      "FLOE f1_metric: 67.405%\n",
      "Combined score: 71.793%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.011938333511352539,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d93b7420668f4db1baef52c11ad66308",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 2.002\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.014322042465209961,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e2833f5444e498da61f1c3d215c61b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 1.061\n",
      "Epoch 6 score:\n",
      "SIC r2_metric: 56.373%\n",
      "SOD f1_metric: 71.866%\n",
      "FLOE f1_metric: 65.286%\n",
      "Combined score: 64.353%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.012949228286743164,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e0df0892f50044b2bd6cf143baf6d2b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 2.035\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.014992475509643555,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6bac9352c95247ef957c884b771a0bf4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 1.246\n",
      "Epoch 7 score:\n",
      "SIC r2_metric: 64.374%\n",
      "SOD f1_metric: 81.755%\n",
      "FLOE f1_metric: 67.67%\n",
      "Combined score: 71.986%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.012461185455322266,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d0eb12fde2aa48b7862918875a4da538",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 1.905\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.015167951583862305,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e6403ad8784a4d118fd70a1ff0bf25e4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 1.238\n",
      "Epoch 8 score:\n",
      "SIC r2_metric: 60.054%\n",
      "SOD f1_metric: 83.863%\n",
      "FLOE f1_metric: 67.746%\n",
      "Combined score: 71.116%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.012663125991821289,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d3bce044285d4d458c06db693301a81e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 1.903\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.014294624328613281,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "01f76ecd9b01490995fe5fdd7c3f7fa5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 0.516\n",
      "Epoch 9 score:\n",
      "SIC r2_metric: 64.564%\n",
      "SOD f1_metric: 85.115%\n",
      "FLOE f1_metric: 67.259%\n",
      "Combined score: 73.323%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.012660980224609375,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "786429415c9247209f8c0c3704a7c8ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 1.848\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.013560056686401367,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "195eda7ff56541e0b2b7b6a84faeb0b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 0.604\n",
      "Epoch 10 score:\n",
      "SIC r2_metric: 61.097%\n",
      "SOD f1_metric: 85.321%\n",
      "FLOE f1_metric: 67.477%\n",
      "Combined score: 72.063%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.013039588928222656,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d112887c1504fc7974dbaae6f444070",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 1.825\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.013839960098266602,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ca0241202cde404bad21692553e02e45",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 0.950\n",
      "Epoch 11 score:\n",
      "SIC r2_metric: 60.932%\n",
      "SOD f1_metric: 84.501%\n",
      "FLOE f1_metric: 67.628%\n",
      "Combined score: 71.699%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.012641668319702148,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f20a3d9c1c8e4f4dbc1634e843faa3a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean training loss: 1.850\r"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.013828754425048828,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 10,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fec110729db2481ebe73661dfe15a3a5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Final batch loss: 1.034\n",
      "Epoch 12 score:\n",
      "SIC r2_metric: 68.137%\n",
      "SOD f1_metric: 83.279%\n",
      "FLOE f1_metric: 67.671%\n",
      "Combined score: 74.101%\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.012656211853027344,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": 24,
       "postfix": null,
       "prefix": "",
       "rate": null,
       "total": 500,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9f15f165d2d040299c1edb5a7b75c709",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/500 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cropping in 20180922T145543_cis_prep.nc failed.\n",
      "Scene size: (3540, 5226) for crop shape: (256, 256)\n",
      "Skipping scene.\n",
      "Cropping in 20181102T212609_cis_prep.nc failed.\n",
      "Cropping in 20180914T115342_cis_prep.nc failed.Scene size: (4050, 5238) for crop shape: (256, 256)\n",
      "Skipping scene.\n",
      "\n",
      "Scene size: (5000, 5312) for crop shape: (256, 256)\n",
      "Skipping scene.\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Input \u001b[0;32mIn [5]\u001b[0m, in \u001b[0;36m<cell line: 4>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      7\u001b[0m net\u001b[38;5;241m.\u001b[39mtrain()  \u001b[38;5;66;03m# Set network to evaluation mode.\u001b[39;00m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;66;03m# Loops though batches in queue.\u001b[39;00m\n\u001b[0;32m---> 10\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i, (batch_x, batch_y) \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(tqdm(iterable\u001b[38;5;241m=\u001b[39mdataloader, total\u001b[38;5;241m=\u001b[39mtrain_options[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mepoch_len\u001b[39m\u001b[38;5;124m'\u001b[39m], colour\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mred\u001b[39m\u001b[38;5;124m'\u001b[39m, position\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)):\n\u001b[1;32m     11\u001b[0m     torch\u001b[38;5;241m.\u001b[39mcuda\u001b[38;5;241m.\u001b[39mempty_cache()  \u001b[38;5;66;03m# Empties the GPU cache freeing up memory.\u001b[39;00m\n\u001b[1;32m     12\u001b[0m     loss_batch \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m  \u001b[38;5;66;03m# Reset from previous batch.\u001b[39;00m\n",
      "File \u001b[0;32m~/AI4ArcticComp/lib/python3.9/site-packages/tqdm/notebook.py:258\u001b[0m, in \u001b[0;36mtqdm_notebook.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    256\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    257\u001b[0m     it \u001b[38;5;241m=\u001b[39m \u001b[38;5;28msuper\u001b[39m(tqdm_notebook, \u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__iter__\u001b[39m()\n\u001b[0;32m--> 258\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m obj \u001b[38;5;129;01min\u001b[39;00m it:\n\u001b[1;32m    259\u001b[0m         \u001b[38;5;66;03m# return super(tqdm...) will not catch exception\u001b[39;00m\n\u001b[1;32m    260\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m obj\n\u001b[1;32m    261\u001b[0m \u001b[38;5;66;03m# NB: except ... [ as ...] breaks IPython async KeyboardInterrupt\u001b[39;00m\n",
      "File \u001b[0;32m~/AI4ArcticComp/lib/python3.9/site-packages/tqdm/std.py:1195\u001b[0m, in \u001b[0;36mtqdm.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1192\u001b[0m time \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_time\n\u001b[1;32m   1194\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1195\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m obj \u001b[38;5;129;01min\u001b[39;00m iterable:\n\u001b[1;32m   1196\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m obj\n\u001b[1;32m   1197\u001b[0m         \u001b[38;5;66;03m# Update and possibly print the progressbar.\u001b[39;00m\n\u001b[1;32m   1198\u001b[0m         \u001b[38;5;66;03m# Note: does not call self.update(1) for speed optimisation.\u001b[39;00m\n",
      "File \u001b[0;32m~/AI4ArcticComp/lib/python3.9/site-packages/torch/utils/data/dataloader.py:681\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    678\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    679\u001b[0m     \u001b[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[39;00m\n\u001b[1;32m    680\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()  \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[0;32m--> 681\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_next_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    682\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    683\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    684\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    685\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called:\n",
      "File \u001b[0;32m~/AI4ArcticComp/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1359\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1356\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_process_data(data)\n\u001b[1;32m   1358\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_shutdown \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_tasks_outstanding \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[0;32m-> 1359\u001b[0m idx, data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_get_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1360\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_tasks_outstanding \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m   1361\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable:\n\u001b[1;32m   1362\u001b[0m     \u001b[38;5;66;03m# Check for _IterableDatasetStopIteration\u001b[39;00m\n",
      "File \u001b[0;32m~/AI4ArcticComp/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1315\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._get_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1313\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory:\n\u001b[1;32m   1314\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory_thread\u001b[38;5;241m.\u001b[39mis_alive():\n\u001b[0;32m-> 1315\u001b[0m         success, data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_try_get_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1316\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m success:\n\u001b[1;32m   1317\u001b[0m             \u001b[38;5;28;01mreturn\u001b[39;00m data\n",
      "File \u001b[0;32m~/AI4ArcticComp/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1163\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m   1150\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_try_get_data\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout\u001b[38;5;241m=\u001b[39m_utils\u001b[38;5;241m.\u001b[39mMP_STATUS_CHECK_INTERVAL):\n\u001b[1;32m   1151\u001b[0m     \u001b[38;5;66;03m# Tries to fetch data from `self._data_queue` once for a given timeout.\u001b[39;00m\n\u001b[1;32m   1152\u001b[0m     \u001b[38;5;66;03m# This can also be used as inner loop of fetching without timeout, with\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1160\u001b[0m     \u001b[38;5;66;03m# Returns a 2-tuple:\u001b[39;00m\n\u001b[1;32m   1161\u001b[0m     \u001b[38;5;66;03m#   (bool: whether successfully get data, any: data if successful else None)\u001b[39;00m\n\u001b[1;32m   1162\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1163\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_data_queue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1164\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m (\u001b[38;5;28;01mTrue\u001b[39;00m, data)\n\u001b[1;32m   1165\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m   1166\u001b[0m         \u001b[38;5;66;03m# At timeout and error, we manually check whether any worker has\u001b[39;00m\n\u001b[1;32m   1167\u001b[0m         \u001b[38;5;66;03m# failed. Note that this is the only mechanism for Windows to detect\u001b[39;00m\n\u001b[1;32m   1168\u001b[0m         \u001b[38;5;66;03m# worker failures.\u001b[39;00m\n",
      "File \u001b[0;32m/appl/python/3.9.11/lib/python3.9/queue.py:180\u001b[0m, in \u001b[0;36mQueue.get\u001b[0;34m(self, block, timeout)\u001b[0m\n\u001b[1;32m    178\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m remaining \u001b[38;5;241m<\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0.0\u001b[39m:\n\u001b[1;32m    179\u001b[0m             \u001b[38;5;28;01mraise\u001b[39;00m Empty\n\u001b[0;32m--> 180\u001b[0m         \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnot_empty\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwait\u001b[49m\u001b[43m(\u001b[49m\u001b[43mremaining\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    181\u001b[0m item \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get()\n\u001b[1;32m    182\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnot_full\u001b[38;5;241m.\u001b[39mnotify()\n",
      "File \u001b[0;32m/appl/python/3.9.11/lib/python3.9/threading.py:316\u001b[0m, in \u001b[0;36mCondition.wait\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    314\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    315\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m timeout \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m--> 316\u001b[0m         gotit \u001b[38;5;241m=\u001b[39m \u001b[43mwaiter\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43macquire\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    317\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    318\u001b[0m         gotit \u001b[38;5;241m=\u001b[39m waiter\u001b[38;5;241m.\u001b[39macquire(\u001b[38;5;28;01mFalse\u001b[39;00m)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "best_combined_score = 0  # Best weighted model score.\n",
    "\n",
    "# -- Training Loop -- #\n",
    "for epoch in tqdm(iterable=range(train_options['epochs']), position=0):\n",
    "    gc.collect()  # Collect garbage to free memory.\n",
    "    loss_sum = torch.tensor([0.])  # To sum the batch losses during the epoch.\n",
    "    net.train()  # Set network to evaluation mode.\n",
    "\n",
    "    # Loops though batches in queue.\n",
    "    for i, (batch_x, batch_y) in enumerate(tqdm(iterable=dataloader, total=train_options['epoch_len'], colour='red', position=0)):\n",
    "        torch.cuda.empty_cache()  # Empties the GPU cache freeing up memory.\n",
    "        loss_batch = 0  # Reset from previous batch.\n",
    "        \n",
    "        # - Transfer to device.\n",
    "        batch_x = batch_x.to(device, non_blocking=True)\n",
    "\n",
    "        # - Mixed precision training. (Saving memory)\n",
    "        with torch.cuda.amp.autocast():\n",
    "            # - Forward pass. \n",
    "            output = net(batch_x)\n",
    "\n",
    "            # - Calculate loss.\n",
    "            for chart in train_options['charts']:\n",
    "                loss_batch += loss_functions[chart](input=output[chart], target=batch_y[chart].to(device))\n",
    "\n",
    "        # - Reset gradients from previous pass.\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # - Backward pass.\n",
    "        loss_batch.backward()\n",
    "\n",
    "        # - Optimizer step\n",
    "        optimizer.step()\n",
    "\n",
    "        # - Add batch loss.\n",
    "        loss_sum += loss_batch.detach().item()\n",
    "\n",
    "        # - Average loss for displaying\n",
    "        loss_epoch = torch.true_divide(loss_sum, i + 1).detach().item()\n",
    "        print('\\rMean training loss: ' + f'{loss_epoch:.3f}', end='\\r')\n",
    "        del output, batch_x, batch_y # Free memory.\n",
    "    del loss_sum\n",
    "\n",
    "    # -- Validation Loop -- #\n",
    "    loss_batch = loss_batch.detach().item()  # For printing after the validation loop.\n",
    "    \n",
    "    # - Stores the output and the reference pixels to calculate the scores after inference on all the scenes.\n",
    "    outputs_flat = {chart: np.array([]) for chart in train_options['charts']}\n",
    "    inf_ys_flat = {chart: np.array([]) for chart in train_options['charts']}\n",
    "\n",
    "    net.eval()  # Set network to evaluation mode.\n",
    "    # - Loops though scenes in queue.\n",
    "    for inf_x, inf_y, masks, name in tqdm(iterable=dataloader_val, total=len(train_options['validate_list']), colour='green', position=0):\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "        # - Ensures that no gradients are calculated, which otherwise take up a lot of space on the GPU.\n",
    "        with torch.no_grad(), torch.cuda.amp.autocast():\n",
    "            inf_x = inf_x.to(device, non_blocking=True)\n",
    "            output = net(inf_x)\n",
    "    \n",
    "        # - Final output layer, and storing of non masked pixels.\n",
    "        for chart in train_options['charts']:\n",
    "            output[chart] = torch.argmax(output[chart], dim=1).squeeze().cpu().numpy()\n",
    "            outputs_flat[chart] = np.append(outputs_flat[chart], output[chart][~masks[chart]])\n",
    "            inf_ys_flat[chart] = np.append(inf_ys_flat[chart], inf_y[chart][~masks[chart]].numpy())\n",
    "        \n",
    "        del inf_x, inf_y, masks, output  # Free memory.\n",
    "\n",
    "    # - Compute the relevant scores.\n",
    "    combined_score, scores = compute_metrics(true=inf_ys_flat, pred=outputs_flat, charts=train_options['charts'],\n",
    "                                             metrics=train_options['chart_metric'])\n",
    "\n",
    "    print(\"\")\n",
    "    print(f\"Final batch loss: {loss_batch:.3f}\")\n",
    "    print(f\"Epoch {epoch} score:\")\n",
    "    for chart in train_options['charts']:\n",
    "        print(f\"{chart} {train_options['chart_metric'][chart]['func'].__name__}: {scores[chart]}%\")\n",
    "    print(f\"Combined score: {combined_score}%\")\n",
    "\n",
    "    # If the scores is better than the previous epoch, then save the model and rename the image to best_validation.\n",
    "    if combined_score > best_combined_score:\n",
    "        best_combined_score = combined_score\n",
    "        torch.save(obj={'model_state_dict': net.state_dict(),\n",
    "                        'optimizer_state_dict': optimizer.state_dict(),\n",
    "                        'epoch': epoch},\n",
    "                        f='best_model')\n",
    "    del inf_ys_flat, outputs_flat  # Free memory.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}