Train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tbu_ucRSo5zT"
   },
   "source": [
    "The following is an example of how to utilize our Sen1Floods11 dataset for training a FCNN. In this example, we train and validate on hand-labeled chips of flood events. However, our dataset includes several other options that are detailed in the README. To replace the dataset, as outlined further below, simply replace the train, test, and validation split csv's, and download the corresponding dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9TQtMrI_VhKk"
   },
   "source": [
    "Authenticate Google Cloud Platform. Note that to run this code, you must connect your notebook runtime to a GPU. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qCEt8eNtU9Zm"
   },
   "outputs": [],
   "source": [
    "from google.colab import auth\n",
    "auth.authenticate_user()\n",
    "\n",
    "!curl https://sdk.cloud.google.com | bash\n",
    "\n",
    "!gcloud init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YkUEnwXQVy4k"
   },
   "outputs": [],
   "source": [
    "!echo \"deb http://packages.cloud.google.com/apt gcsfuse-bionic main\" > /etc/apt/sources.list.d/gcsfuse.list\n",
    "!curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -\n",
    "!apt -qq update\n",
    "!apt -qq install gcsfuse"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vXGTA6vHVyJX"
   },
   "source": [
    "Install RasterIO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NLlVutLzV_pZ"
   },
   "outputs": [],
   "source": [
    "!pip install rasterio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hLqL9C2Rg6eB"
   },
   "source": [
    "Define a model checkpoint folder, for storing network checkpoints during training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yLlIhE-Hg-Ym"
   },
   "outputs": [],
   "source": [
    "%cd /home\n",
    "!sudo mkdir checkpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mwrDM4AjVnbU"
   },
   "source": [
    "Download train, test, and validation splits for both flood water. To download different train, test, and validation splits, simply replace these paths with the path to a csv containing the desired splits. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RFLsGwdRWuO4"
   },
   "outputs": [],
   "source": [
    "!gsutil cp gs://sen1floods11/v1.1/splits/flood_handlabeled/flood_train_data.csv .\n",
    "!gsutil cp gs://sen1floods11/v1.1/splits/flood_handlabeled/flood_test_data.csv .\n",
    "!gsutil cp gs://sen1floods11/v1.1/splits/flood_handlabeled/flood_valid_data.csv ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZCAXpuKVW3eV"
   },
   "source": [
    "Download raw train, test, and validation data. In this example, we are downloading train, test, and validation data of flood images which are hand labeled. However, you can simply replace these paths with whichever dataset you would like to use - further documentation of the Sen1Floods11 dataset and organization is available in the README."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ahAWnrSFW53S"
   },
   "outputs": [],
   "source": [
    "!sudo mkdir files\n",
    "!sudo mkdir files/S1\n",
    "!sudo mkdir files/Labels\n",
    "\n",
    "!gsutil -m rsync -r gs://sen1floods11/v1.1/data/flood_events/HandLabeled/S1Hand files/S1\n",
    "!gsutil -m rsync -r gs://sen1floods11/v1.1/data/flood_events/HandLabeled/LabelHand files/Labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_46CazV3XSCD"
   },
   "source": [
    "Define model training hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fNYQywdWXeLM"
   },
   "outputs": [],
   "source": [
    "LR = 5e-4\n",
    "EPOCHS = 100\n",
    "EPOCHS_PER_UPDATE = 1\n",
    "RUNNAME = \"Sen1Floods11\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W9FJmTnZXjxj"
   },
   "source": [
    "Define functions to process and augment training and testing images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mBkfav0Eajqg"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import transforms\n",
    "import torchvision.transforms.functional as F\n",
    "import random\n",
    "from PIL import Image\n",
    "\n",
    "class InMemoryDataset(torch.utils.data.Dataset):\n",
    "  \n",
    "  def __init__(self, data_list, preprocess_func):\n",
    "    self.data_list = data_list\n",
    "    self.preprocess_func = preprocess_func\n",
    "  \n",
    "  def __getitem__(self, i):\n",
    "    return self.preprocess_func(self.data_list[i])\n",
    "  \n",
    "  def __len__(self):\n",
    "    return len(self.data_list)\n",
    "\n",
    "\n",
    "def processAndAugment(data):\n",
    "  (x,y) = data\n",
    "  im,label = x.copy(), y.copy()\n",
    "\n",
    "  # convert to PIL for easier transforms\n",
    "  im1 = Image.fromarray(im[0])\n",
    "  im2 = Image.fromarray(im[1])\n",
    "  label = Image.fromarray(label.squeeze())\n",
    "\n",
    "  # Get params for random transforms\n",
    "  i, j, h, w = transforms.RandomCrop.get_params(im1, (256, 256))\n",
    "  \n",
    "  im1 = F.crop(im1, i, j, h, w)\n",
    "  im2 = F.crop(im2, i, j, h, w)\n",
    "  label = F.crop(label, i, j, h, w)\n",
    "  if random.random() > 0.5:\n",
    "    im1 = F.hflip(im1)\n",
    "    im2 = F.hflip(im2)\n",
    "    label = F.hflip(label)\n",
    "  if random.random() > 0.5:\n",
    "    im1 = F.vflip(im1)\n",
    "    im2 = F.vflip(im2)\n",
    "    label = F.vflip(label)\n",
    "  \n",
    "  norm = transforms.Normalize([0.6851, 0.5235], [0.0820, 0.1102])\n",
    "  im = torch.stack([transforms.ToTensor()(im1).squeeze(), transforms.ToTensor()(im2).squeeze()])\n",
    "  im = norm(im)\n",
    "  label = transforms.ToTensor()(label).squeeze()\n",
    "  if torch.sum(label.gt(.003) * label.lt(.004)):\n",
    "    label *= 255\n",
    "  label = label.round()\n",
    "\n",
    "  return im, label\n",
    "\n",
    "\n",
    "def processTestIm(data):\n",
    "  (x,y) = data\n",
    "  im,label = x.copy(), y.copy()\n",
    "  norm = transforms.Normalize([0.6851, 0.5235], [0.0820, 0.1102])\n",
    "\n",
    "  # convert to PIL for easier transforms\n",
    "  im_c1 = Image.fromarray(im[0]).resize((512,512))\n",
    "  im_c2 = Image.fromarray(im[1]).resize((512,512))\n",
    "  label = Image.fromarray(label.squeeze()).resize((512,512))\n",
    "\n",
    "  im_c1s = [F.crop(im_c1, 0, 0, 256, 256), F.crop(im_c1, 0, 256, 256, 256),\n",
    "            F.crop(im_c1, 256, 0, 256, 256), F.crop(im_c1, 256, 256, 256, 256)]\n",
    "  im_c2s = [F.crop(im_c2, 0, 0, 256, 256), F.crop(im_c2, 0, 256, 256, 256),\n",
    "            F.crop(im_c2, 256, 0, 256, 256), F.crop(im_c2, 256, 256, 256, 256)]\n",
    "  labels = [F.crop(label, 0, 0, 256, 256), F.crop(label, 0, 256, 256, 256),\n",
    "            F.crop(label, 256, 0, 256, 256), F.crop(label, 256, 256, 256, 256)]\n",
    "\n",
    "  ims = [torch.stack((transforms.ToTensor()(x).squeeze(),\n",
    "                    transforms.ToTensor()(y).squeeze()))\n",
    "                    for (x,y) in zip(im_c1s, im_c2s)]\n",
    "  \n",
    "  ims = [norm(im) for im in ims]\n",
    "  ims = torch.stack(ims)\n",
    "  \n",
    "  labels = [(transforms.ToTensor()(label).squeeze()) for label in labels]\n",
    "  labels = torch.stack(labels)\n",
    "  \n",
    "  if torch.sum(labels.gt(.003) * labels.lt(.004)):\n",
    "    labels *= 255\n",
    "  labels = labels.round()\n",
    "  \n",
    "  return ims, labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uzmZIRuoeAuJ"
   },
   "source": [
    "Load *flood water* train, test, and validation data from splits. In this example, this is the data we will use to train our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rQUnYCIBeG21"
   },
   "outputs": [],
   "source": [
    "from time import time\n",
    "import csv\n",
    "import os\n",
    "import numpy as np\n",
    "import rasterio\n",
    "\n",
    "def getArrFlood(fname):\n",
    "  return rasterio.open(fname).read()\n",
    "\n",
    "def download_flood_water_data_from_list(l):\n",
    "  i = 0\n",
    "  tot_nan = 0\n",
    "  tot_good = 0\n",
    "  flood_data = []\n",
    "  for (im_fname, mask_fname) in l:\n",
    "    if not os.path.exists(os.path.join(\"files/\", im_fname)):\n",
    "      continue\n",
    "    arr_x = np.nan_to_num(getArrFlood(os.path.join(\"files/\", im_fname)))\n",
    "    arr_y = getArrFlood(os.path.join(\"files/\", mask_fname))\n",
    "    arr_y[arr_y == -1] = 255 \n",
    "    \n",
    "    arr_x = np.clip(arr_x, -50, 1)\n",
    "    arr_x = (arr_x + 50) / 51\n",
    "      \n",
    "    if i % 100 == 0:\n",
    "      print(im_fname, mask_fname)\n",
    "    i += 1\n",
    "    flood_data.append((arr_x,arr_y))\n",
    "\n",
    "  return flood_data\n",
    "\n",
    "def load_flood_train_data(input_root, label_root):\n",
    "  fname = \"flood_train_data.csv\"\n",
    "  training_files = []\n",
    "  with open(fname) as f:\n",
    "    for line in csv.reader(f):\n",
    "      training_files.append(tuple((input_root+line[0], label_root+line[1])))\n",
    "\n",
    "  return download_flood_water_data_from_list(training_files)\n",
    "\n",
    "def load_flood_valid_data(input_root, label_root):\n",
    "  fname = \"flood_valid_data.csv\"\n",
    "  validation_files = []\n",
    "  with open(fname) as f:\n",
    "    for line in csv.reader(f):\n",
    "      validation_files.append(tuple((input_root+line[0], label_root+line[1])))\n",
    "\n",
    "  return download_flood_water_data_from_list(validation_files)\n",
    "\n",
    "def load_flood_test_data(input_root, label_root):\n",
    "  fname = \"flood_test_data.csv\"\n",
    "  testing_files = []\n",
    "  with open(fname) as f:\n",
    "    for line in csv.reader(f):\n",
    "      testing_files.append(tuple((input_root+line[0], label_root+line[1])))\n",
    "  \n",
    "  return download_flood_water_data_from_list(testing_files)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cFp9jrHYfOUh"
   },
   "source": [
    "Load training data and validation data. Note that here, we have chosen to train and validate our model on flood data. However, you can simply replace the load function call with one of the options defined above to load a different dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZcqPlsjBffXx"
   },
   "outputs": [],
   "source": [
    "train_data = load_flood_train_data('S1/', 'Labels/')\n",
    "train_dataset = InMemoryDataset(train_data, processAndAugment)\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True, sampler=None,\n",
    "                  batch_sampler=None, num_workers=0, collate_fn=None,\n",
    "                  pin_memory=True, drop_last=False, timeout=0,\n",
    "                  worker_init_fn=None)\n",
    "train_iter = iter(train_loader)\n",
    "\n",
    "valid_data = load_flood_valid_data('S1/', 'Labels/') \n",
    "valid_dataset = InMemoryDataset(valid_data, processTestIm)\n",
    "valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=4, shuffle=True, sampler=None,\n",
    "                  batch_sampler=None, num_workers=0, collate_fn=lambda x: (torch.cat([a[0] for a in x], 0), torch.cat([a[1] for a in x], 0)),\n",
    "                  pin_memory=True, drop_last=False, timeout=0,\n",
    "                  worker_init_fn=None)\n",
    "valid_iter = iter(valid_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i3aAhUi2fp7M"
   },
   "source": [
    "Define the network. For our purposes, we use ResNet50. However, if you wish to test a different model framework, optimizer, or loss function you can simply replace those here. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5cp4uXI1f9dr"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision.models as models\n",
    "import torch.nn as nn\n",
    "\n",
    "net = models.segmentation.fcn_resnet50(pretrained=False, num_classes=2, pretrained_backbone=False)\n",
    "net.backbone.conv1 = nn.Conv2d(2, 64, kernel_size=7, stride=2, padding=3, bias=False)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(weight=torch.tensor([1,8]).float().cuda(), ignore_index=255) \n",
    "optimizer = torch.optim.AdamW(net.parameters(),lr=LR)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, len(train_loader) * 10, T_mult=2, eta_min=0, last_epoch=-1)\n",
    "\n",
    "def convertBNtoGN(module, num_groups=16):\n",
    "  if isinstance(module, torch.nn.modules.batchnorm.BatchNorm2d):\n",
    "    return nn.GroupNorm(num_groups, module.num_features,\n",
    "                        eps=module.eps, affine=module.affine)\n",
    "    if module.affine:\n",
    "        mod.weight.data = module.weight.data.clone().detach()\n",
    "        mod.bias.data = module.bias.data.clone().detach()\n",
    "\n",
    "  for name, child in module.named_children():\n",
    "      module.add_module(name, convertBNtoGN(child, num_groups=num_groups))\n",
    "\n",
    "  return module\n",
    "\n",
    "net = convertBNtoGN(net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g_Sy3ALGgQjf"
   },
   "source": [
    "Define assessment metrics. For our purposes, we use overall accuracy and mean intersection over union. However, we also include functions for calculating true positives, false positives, true negatives, and false negatives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bwxC-fVBgUIb"
   },
   "outputs": [],
   "source": [
    "def computeIOU(output, target):\n",
    "  output = torch.argmax(output, dim=1).flatten() \n",
    "  target = target.flatten()\n",
    "  \n",
    "  no_ignore = target.ne(255).cuda()\n",
    "  output = output.masked_select(no_ignore)\n",
    "  target = target.masked_select(no_ignore)\n",
    "  intersection = torch.sum(output * target)\n",
    "  union = torch.sum(target) + torch.sum(output) - intersection\n",
    "  iou = (intersection + .0000001) / (union + .0000001)\n",
    "  \n",
    "  if iou != iou:\n",
    "    print(\"failed, replacing with 0\")\n",
    "    iou = torch.tensor(0).float()\n",
    "  \n",
    "  return iou\n",
    "  \n",
    "def computeAccuracy(output, target):\n",
    "  output = torch.argmax(output, dim=1).flatten() \n",
    "  target = target.flatten()\n",
    "  \n",
    "  no_ignore = target.ne(255).cuda()\n",
    "  output = output.masked_select(no_ignore)\n",
    "  target = target.masked_select(no_ignore)\n",
    "  correct = torch.sum(output.eq(target))\n",
    "  \n",
    "  return correct.float() / len(target)\n",
    "\n",
    "def truePositives(output, target):\n",
    "  output = torch.argmax(output, dim=1).flatten() \n",
    "  target = target.flatten()\n",
    "  no_ignore = target.ne(255).cuda()\n",
    "  output = output.masked_select(no_ignore)\n",
    "  target = target.masked_select(no_ignore)\n",
    "  correct = torch.sum(output * target)\n",
    "  \n",
    "  return correct\n",
    "\n",
    "def trueNegatives(output, target):\n",
    "  output = torch.argmax(output, dim=1).flatten() \n",
    "  target = target.flatten()\n",
    "  no_ignore = target.ne(255).cuda()\n",
    "  output = output.masked_select(no_ignore)\n",
    "  target = target.masked_select(no_ignore)\n",
    "  output = (output == 0)\n",
    "  target = (target == 0)\n",
    "  correct = torch.sum(output * target)\n",
    "  \n",
    "  return correct\n",
    "\n",
    "def falsePositives(output, target):\n",
    "  output = torch.argmax(output, dim=1).flatten() \n",
    "  target = target.flatten()\n",
    "  no_ignore = target.ne(255).cuda()\n",
    "  output = output.masked_select(no_ignore)\n",
    "  target = target.masked_select(no_ignore)\n",
    "  output = (output == 1)\n",
    "  target = (target == 0)\n",
    "  correct = torch.sum(output * target)\n",
    "  \n",
    "  return correct\n",
    "\n",
    "def falseNegatives(output, target):\n",
    "  output = torch.argmax(output, dim=1).flatten() \n",
    "  target = target.flatten()\n",
    "  no_ignore = target.ne(255).cuda()\n",
    "  output = output.masked_select(no_ignore)\n",
    "  target = target.masked_select(no_ignore)\n",
    "  output = (output == 0)\n",
    "  target = (target == 1)\n",
    "  correct = torch.sum(output * target)\n",
    "  \n",
    "  return correct"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lun5tGoYgjWX"
   },
   "source": [
    "Define training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DubsYZ8GgkxD"
   },
   "outputs": [],
   "source": [
    "training_losses = []\n",
    "training_accuracies = []\n",
    "training_ious = []\n",
    "\n",
    "def train_loop(inputs, labels, net, optimizer, scheduler):\n",
    "  global running_loss\n",
    "  global running_iou\n",
    "  global running_count\n",
    "  global running_accuracy\n",
    "  \n",
    "  # zero the parameter gradients\n",
    "  optimizer.zero_grad()\n",
    "  net = net.cuda()\n",
    "  \n",
    "  # forward + backward + optimize\n",
    "  outputs = net(inputs.cuda())\n",
    "  loss = criterion(outputs[\"out\"], labels.long().cuda())\n",
    "  loss.backward()\n",
    "  optimizer.step()\n",
    "  scheduler.step()\n",
    "\n",
    "  running_loss += loss\n",
    "  running_iou += computeIOU(outputs[\"out\"], labels.cuda())\n",
    "  running_accuracy += computeAccuracy(outputs[\"out\"], labels.cuda())\n",
    "  running_count += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iM3Jz__hgshh"
   },
   "source": [
    "Define validation loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_GmVaoRvguic"
   },
   "outputs": [],
   "source": [
    "valid_losses = []\n",
    "valid_accuracies = []\n",
    "valid_ious = []\n",
    "\n",
    "def validation_loop(validation_data_loader, net):\n",
    "  global running_loss\n",
    "  global running_iou\n",
    "  global running_count\n",
    "  global running_accuracy\n",
    "  global max_valid_iou\n",
    "\n",
    "  global training_losses\n",
    "  global training_accuracies\n",
    "  global training_ious\n",
    "  global valid_losses\n",
    "  global valid_accuracies\n",
    "  global valid_ious\n",
    "\n",
    "  net = net.eval()\n",
    "  net = net.cuda()\n",
    "  count = 0\n",
    "  iou = 0\n",
    "  loss = 0\n",
    "  accuracy = 0\n",
    "  with torch.no_grad():\n",
    "      for (images, labels) in validation_data_loader:\n",
    "          net = net.cuda()\n",
    "          outputs = net(images.cuda())\n",
    "          valid_loss = criterion(outputs[\"out\"], labels.long().cuda())\n",
    "          valid_iou = computeIOU(outputs[\"out\"], labels.cuda())\n",
    "          valid_accuracy = computeAccuracy(outputs[\"out\"], labels.cuda())\n",
    "          iou += valid_iou\n",
    "          loss += valid_loss\n",
    "          accuracy += valid_accuracy\n",
    "          count += 1\n",
    "\n",
    "  iou = iou / count\n",
    "  accuracy = accuracy / count\n",
    "\n",
    "  if iou > max_valid_iou:\n",
    "    max_valid_iou = iou\n",
    "    save_path = os.path.join(\"checkpoints\", \"{}_{}_{}.cp\".format(RUNNAME, i, iou.item()))\n",
    "    torch.save(net.state_dict(), save_path)\n",
    "    print(\"model saved at\", save_path)\n",
    "\n",
    "  loss = loss / count\n",
    "  print(\"Training Loss:\", running_loss / running_count)\n",
    "  print(\"Training IOU:\", running_iou / running_count)\n",
    "  print(\"Training Accuracy:\", running_accuracy / running_count)\n",
    "  print(\"Validation Loss:\", loss)\n",
    "  print(\"Validation IOU:\", iou)\n",
    "  print(\"Validation Accuracy:\", accuracy)\n",
    "\n",
    "\n",
    "  training_losses.append(running_loss / running_count)\n",
    "  training_accuracies.append(running_accuracy / running_count)\n",
    "  training_ious.append(running_iou / running_count)\n",
    "  valid_losses.append(loss)\n",
    "  valid_accuracies.append(accuracy)\n",
    "  valid_ious.append(iou)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DBMattYshiUj"
   },
   "source": [
    "Define testing loop (here, you can replace assessment metrics)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mI_mhL_ehjot"
   },
   "outputs": [],
   "source": [
    "def test_loop(test_data_loader, net):\n",
    "  net = net.eval()\n",
    "  net = net.cuda()\n",
    "  count = 0\n",
    "  iou = 0\n",
    "  loss = 0\n",
    "  accuracy = 0\n",
    "  with torch.no_grad():\n",
    "      for (images, labels) in tqdm(test_data_loader):\n",
    "          net = net.cuda()\n",
    "          outputs = net(images.cuda())\n",
    "          valid_loss = criterion(outputs[\"out\"], labels.long().cuda())\n",
    "          valid_iou = computeIOU(outputs[\"out\"], labels.cuda())\n",
    "          iou += valid_iou\n",
    "          accuracy += computeAccuracy(outputs[\"out\"], labels.cuda())\n",
    "          count += 1\n",
    "\n",
    "  iou = iou / count\n",
    "  print(\"Test IOU:\", iou)\n",
    "  print(\"Test Accuracy:\", accuracy / count)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cy9Fii06h17Q"
   },
   "source": [
    "Define training and validation scheme"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NZuKVC6wh4Go"
   },
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "from IPython.display import clear_output\n",
    "\n",
    "running_loss = 0\n",
    "running_iou = 0\n",
    "running_count = 0\n",
    "running_accuracy = 0\n",
    "\n",
    "training_losses = []\n",
    "training_accuracies = []\n",
    "training_ious = []\n",
    "valid_losses = []\n",
    "valid_accuracies = []\n",
    "valid_ious = []\n",
    "\n",
    "\n",
    "def train_epoch(net, optimizer, scheduler, train_iter):\n",
    "  for (inputs, labels) in tqdm(train_iter):\n",
    "    train_loop(inputs.cuda(), labels.cuda(), net.cuda(), optimizer, scheduler)\n",
    " \n",
    "\n",
    "def train_validation_loop(net, optimizer, scheduler, train_loader,\n",
    "                          valid_loader, num_epochs, cur_epoch):\n",
    "  global running_loss\n",
    "  global running_iou\n",
    "  global running_count\n",
    "  global running_accuracy\n",
    "  net = net.train()\n",
    "  running_loss = 0\n",
    "  running_iou = 0\n",
    "  running_count = 0\n",
    "  running_accuracy = 0\n",
    "  \n",
    "  for i in tqdm(range(num_epochs)):\n",
    "    train_iter = iter(train_loader)\n",
    "    train_epoch(net, optimizer, scheduler, train_iter)\n",
    "  clear_output()\n",
    "  \n",
    "  print(\"Current Epoch:\", cur_epoch)\n",
    "  validation_loop(iter(valid_loader), net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k3I88aY5iAWD"
   },
   "source": [
    "Train model and assess metrics over epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8MRpxUGWiDTu"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from IPython.display import display\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "max_valid_iou = 0\n",
    "start = 0\n",
    "\n",
    "epochs = []\n",
    "training_losses = []\n",
    "training_accuracies = []\n",
    "training_ious = []\n",
    "valid_losses = []\n",
    "valid_accuracies = []\n",
    "valid_ious = []\n",
    "\n",
    "for i in range(start, 1000):\n",
    "  train_validation_loop(net, optimizer, scheduler, train_loader, valid_loader, 10, i)\n",
    "  epochs.append(i)\n",
    "  x = epochs\n",
    "  plt.plot(x, training_losses, label='training losses')\n",
    "  plt.plot(x, training_accuracies, 'tab:orange', label='training accuracy')\n",
    "  plt.plot(x, training_ious, 'tab:purple', label='training iou')\n",
    "  plt.plot(x, valid_losses, label='valid losses')\n",
    "  plt.plot(x, valid_accuracies, 'tab:red',label='valid accuracy')\n",
    "  plt.plot(x, valid_ious, 'tab:green',label='valid iou')\n",
    "  plt.legend(loc=\"upper left\")\n",
    "\n",
    "  display(plt.show())\n",
    "\n",
    "  print(\"max valid iou:\", max_valid_iou)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Train.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}