combination_trainer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#clear all variables\n",
    "%reset -f\n",
    "# from class_and_functions_for_combinations import *\n",
    "from class_and_functions_for_combinations import prepare_ds, my_load, MyDataset, HEstimator, reset_weights, DataLoader, train_epoch, val_epoch, deepcopy, train, evaluate\n",
    "from class_and_functions_for_combinations import ALL_EVALUATION_DATASETS, get_best_result, REAL_CLEAN_DATASETS, REAL_NOISY_DATASETS, get_STDs_for, get2D_MSEs_for, ALL_SIM_EVAL_DATASETS, ALL_REAL_EVAL_DATASETS\n",
    "import onnx, numpy as np, os, torch, torchvision, time, copy, sys, pickle, cv2 as cv\n",
    "from tqdm import tqdm\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "# device = torch.device(\"cuda:2\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f'device: {device}')\n",
    "dir_names = ['tmp', 'tmp/dss', 'tmp/evals', 'tmp/hes', 'tmp/models', 'tmp/real_dss/', 'tmp/training_combinations/']\n",
    "for dir_name in dir_names:\n",
    "    if not os.path.exists(dir_name):\n",
    "        print('creating directory: ', dir_name)\n",
    "        os.makedirs(dir_name)\n",
    "# device = torch.device(\"cpu\")\n",
    "#clear cuda memory\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Param Combinations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #### STEER NOISE VS HE _DISTANCE ####\n",
    "# architecture_vars = ['a']\n",
    "# droput_vars = [.4]#[0,.4,.8]\n",
    "\n",
    "# steer_noise_level_vars = [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20] #[8, 12]#[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]#[0,4,8,10,12,16,20]#[0,2, 4,6, 8, 10, 12, 14, 16, 18, 20]\n",
    "# he_distance_vars = [.2,.3,.4,.5,.6,.7,.8,.9] #[.5,.6]#[.3,.4,.5,.6,.7,.8,.9]#[.2,.3,.4,.5,.6,.7,.8,.9]#[.4,.6,.8]#[.3,.4,.5,.55,.6,.7,.8,.9]\n",
    "\n",
    "# canny1_vars = [100]#[100, 0]\n",
    "# canny2_vars = [200]#[200, 0]\n",
    "# blur_vars = [3]#[0,3,5,7,9]#[3]#[0,3,5,7,9]\n",
    "# img_noise_vars = [80]#[0,40,80,120,160]#[80]#[0,40,80,120,160]\n",
    "# keep_bottom_vars = [.8]#[.8]#[.4,.5,.6,.7,.8,.9]#[.8]#[.4,.5,.6,.7,.8,.9]\n",
    "# ds_length_vars = [10000]\n",
    "# img_size_vars = [32]\n",
    "\n",
    "# batch_size_vars = [2**16]#[2**8, 2**9,2**10,2**11,2**12,2**13,2**14,2**15,2**16]\n",
    "# lr_vars =  [3e-3, 5e-3, 3e-3]#[1e-1, 5e-2, 1e-2, 5e-3, 3e-3, 1e-3, 1e-4, 1e-5]#[0.003]\n",
    "# epochs_vars = [100]#[50,100,200, 300, 500]#[50, 100, 200]\n",
    "# L1_lambda_vars = [1e-3]#[1e-4]#[1e-2,1e-3,1e-4,1e-5,1e-6]\n",
    "# L2_lambda_vars = [1e-2]#[1e-1,1e-2,1e-3,1e-4,1e-5]\n",
    "# weight_decay_vars = [9e-5]#[9e-2,9e-3,9e-4,9e-5,9e-6]#[9e-5] #5e-2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### REGULARIzation\n",
    "# architecture_vars = ['a']\n",
    "# droput_vars = [.4]#[0,.4,.8]\n",
    "\n",
    "# steer_noise_level_vars = [10, 12, 14] #[8, 12]#[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]#[0,4,8,10,12,16,20]#[0,2, 4,6, 8, 10, 12, 14, 16, 18, 20]\n",
    "# he_distance_vars = [.4,.5] #[.5,.6]#[.3,.4,.5,.6,.7,.8,.9]#[.2,.3,.4,.5,.6,.7,.8,.9]#[.4,.6,.8]#[.3,.4,.5,.55,.6,.7,.8,.9]\n",
    "\n",
    "# canny1_vars = [100]#[100, 0]\n",
    "# canny2_vars = [200]#[200, 0]\n",
    "# blur_vars = [3]#[0,3,5,7,9]#[3]#[0,3,5,7,9]\n",
    "# img_noise_vars = [80]#[0,40,80,120,160]#[80]#[0,40,80,120,160]\n",
    "# keep_bottom_vars = [.8]#[.8]#[.4,.5,.6,.7,.8,.9]#[.8]#[.4,.5,.6,.7,.8,.9]\n",
    "# ds_length_vars = [10000]\n",
    "# img_size_vars = [32]\n",
    "\n",
    "# batch_size_vars = [2**16]#[2**8, 2**9,2**10,2**11,2**12,2**13,2**14,2**15,2**16]\n",
    "# lr_vars =  [5e-3]#[3e-3, 5e-3, 3e-3]#[1e-1, 5e-2, 1e-2, 5e-3, 3e-3, 1e-3, 1e-4, 1e-5]#[0.003]\n",
    "# epochs_vars = [100]#[50,100,200, 300, 500]#[50, 100, 200]\n",
    "# L1_lambda_vars = [1e-3]#[1e-4]#[1e-2,1e-3,1e-4,1e-5,1e-6]\n",
    "# L2_lambda_vars = [1e-1,1e-2,1e-3,1e-4]#[1e-2]#[1e-1,1e-2,1e-3,1e-4,1e-5]\n",
    "# weight_decay_vars = [9e-5]#[9e-2,9e-3,9e-4,9e-5,9e-6]#[9e-5] #5e-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## DROPOUT\n",
    "# architecture_vars = ['a']\n",
    "# droput_vars = [.1,.2,.3,.4,.5,.6,.8]#[0,.4,.8]\n",
    "\n",
    "# steer_noise_level_vars = [10, 12, 14] #[8, 12]#[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]#[0,4,8,10,12,16,20]#[0,2, 4,6, 8, 10, 12, 14, 16, 18, 20]\n",
    "# he_distance_vars = [.4,.5] #[.5,.6]#[.3,.4,.5,.6,.7,.8,.9]#[.2,.3,.4,.5,.6,.7,.8,.9]#[.4,.6,.8]#[.3,.4,.5,.55,.6,.7,.8,.9]\n",
    "\n",
    "# canny1_vars = [100]#[100, 0]\n",
    "# canny2_vars = [200]#[200, 0]\n",
    "# blur_vars = [3]#[0,3,5,7,9]#[3]#[0,3,5,7,9]\n",
    "# img_noise_vars = [80]#[0,40,80,120,160]#[80]#[0,40,80,120,160]\n",
    "# keep_bottom_vars = [.8]#[.8]#[.4,.5,.6,.7,.8,.9]#[.8]#[.4,.5,.6,.7,.8,.9]\n",
    "# ds_length_vars = [10000]\n",
    "# img_size_vars = [32]\n",
    "\n",
    "# batch_size_vars = [2**16]#[2**8, 2**9,2**10,2**11,2**12,2**13,2**14,2**15,2**16]\n",
    "# lr_vars =  [5e-3]#[3e-3, 5e-3, 3e-3]#[1e-1, 5e-2, 1e-2, 5e-3, 3e-3, 1e-3, 1e-4, 1e-5]#[0.003]\n",
    "# epochs_vars = [100]#[50,100,200, 300, 500]#[50, 100, 200]\n",
    "# L1_lambda_vars = [1e-3]#[1e-4]#[1e-2,1e-3,1e-4,1e-5,1e-6]\n",
    "# L2_lambda_vars = [1e-2]#[1e-1,1e-2,1e-3,1e-4,1e-5]\n",
    "# weight_decay_vars = [9e-5]#[9e-2,9e-3,9e-4,9e-5,9e-6]#[9e-5] #5e-2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## LR & EPOCHS\n",
    "# architecture_vars = ['a']\n",
    "# droput_vars = [.4]#[0,.4,.8]\n",
    "\n",
    "# steer_noise_level_vars = [8, 12]#[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]#[0,4,8,10,12,16,20]#[0,2, 4,6, 8, 10, 12, 14, 16, 18, 20]\n",
    "# he_distance_vars = [.5,.6]#[.3,.4,.5,.6,.7,.8,.9]#[.2,.3,.4,.5,.6,.7,.8,.9]#[.4,.6,.8]#[.3,.4,.5,.55,.6,.7,.8,.9]\n",
    "\n",
    "# canny1_vars = [100]#[100, 0]\n",
    "# canny2_vars = [200]#[200, 0]\n",
    "# blur_vars = [3]#[0,3,5,7,9]#[3]#[0,3,5,7,9]\n",
    "# img_noise_vars = [80]#[0,40,80,120,160]#[80]#[0,40,80,120,160]\n",
    "# keep_bottom_vars = [.8]#[.8]#[.4,.5,.6,.7,.8,.9]#[.8]#[.4,.5,.6,.7,.8,.9]\n",
    "# ds_length_vars = [10000]\n",
    "# img_size_vars = [32]\n",
    "\n",
    "# batch_size_vars = [2**16]#[2**8, 2**9,2**10,2**11,2**12,2**13,2**14,2**15,2**16]\n",
    "# lr_vars =  [1e-1, 5e-2, 1e-2, 5e-3, 3e-3, 1e-3, 1e-4, 1e-5]#[0.003]\n",
    "# epochs_vars = [50,100,200, 300, 500]#[50, 100, 200]\n",
    "# L1_lambda_vars = [1e-3]#[1e-4]#[1e-2,1e-3,1e-4,1e-5,1e-6]\n",
    "# L2_lambda_vars = [1e-2]#[1e-1,1e-2,1e-3,1e-4,1e-5]\n",
    "# weight_decay_vars = [9e-5]#[9e-2,9e-3,9e-4,9e-5,9e-6]#[9e-5] #5e-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### KEEP BOTTOM\n",
    "# architecture_vars = ['a']\n",
    "# droput_vars = [.3]#[0,.4,.8]\n",
    "\n",
    "# steer_noise_level_vars = [8, 10, 12, 14] #[8, 12]#[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]#[0,4,8,10,12,16,20]#[0,2, 4,6, 8, 10, 12, 14, 16, 18, 20]\n",
    "# he_distance_vars = [.2,.3,.4,.5,.6,.7,.8] #[.5,.6]#[.3,.4,.5,.6,.7,.8,.9]#[.2,.3,.4,.5,.6,.7,.8,.9]#[.4,.6,.8]#[.3,.4,.5,.55,.6,.7,.8,.9]\n",
    "\n",
    "# canny1_vars = [100]#[100, 0]\n",
    "# canny2_vars = [200]#[200, 0]\n",
    "# blur_vars = [3]#[0,3,5,7,9]#[3]#[0,3,5,7,9]\n",
    "# img_noise_vars = [80]#[0,40,80,120,160]#[80]#[0,40,80,120,160]\n",
    "# keep_bottom_vars = [.4,.5,.6,.7,.8,.9]#[.8]#[.4,.5,.6,.7,.8,.9]#[.8]#[.4,.5,.6,.7,.8,.9]\n",
    "# ds_length_vars = [10000]\n",
    "# img_size_vars = [32]\n",
    "\n",
    "# batch_size_vars = [2**16]#[2**8, 2**9,2**10,2**11,2**12,2**13,2**14,2**15,2**16]\n",
    "# lr_vars =  [5e-3]#[3e-3, 5e-3, 3e-3]#[1e-1, 5e-2, 1e-2, 5e-3, 3e-3, 1e-3, 1e-4, 1e-5]#[0.003]\n",
    "# epochs_vars = [100]#[50,100,200, 300, 500]#[50, 100, 200]\n",
    "# L1_lambda_vars = [1e-3]#[1e-4]#[1e-2,1e-3,1e-4,1e-5,1e-6]\n",
    "# L2_lambda_vars = [1e-2]#[1e-1,1e-2,1e-3,1e-4,1e-5]\n",
    "# weight_decay_vars = [9e-5]#[9e-2,9e-3,9e-4,9e-5,9e-6]#[9e-5] #5e-2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## BLUR\n",
    "# architecture_vars = ['a']\n",
    "# droput_vars = [.3]#[0,.4,.8]\n",
    "\n",
    "# steer_noise_level_vars = [10, 12, 14] #[8, 12]#[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]#[0,4,8,10,12,16,20]#[0,2, 4,6, 8, 10, 12, 14, 16, 18, 20]\n",
    "# he_distance_vars = [.4,.5,.6] #[.5,.6]#[.3,.4,.5,.6,.7,.8,.9]#[.2,.3,.4,.5,.6,.7,.8,.9]#[.4,.6,.8]#[.3,.4,.5,.55,.6,.7,.8,.9]\n",
    "\n",
    "# canny1_vars = [100]#[100, 0]\n",
    "# canny2_vars = [200]#[200, 0]\n",
    "# blur_vars = [0,3,5,7]#[0,3,5,7,9]#[3]#[0,3,5,7,9]\n",
    "# img_noise_vars = [80]#[0,40,80,120,160]#[80]#[0,40,80,120,160]\n",
    "# keep_bottom_vars = [.8]#[.8]#[.4,.5,.6,.7,.8,.9]#[.8]#[.4,.5,.6,.7,.8,.9]\n",
    "# ds_length_vars = [10000]\n",
    "# img_size_vars = [32]\n",
    "\n",
    "# batch_size_vars = [2**16]#[2**8, 2**9,2**10,2**11,2**12,2**13,2**14,2**15,2**16]\n",
    "# lr_vars =  [5e-3]#[3e-3, 5e-3, 3e-3]#[1e-1, 5e-2, 1e-2, 5e-3, 3e-3, 1e-3, 1e-4, 1e-5]#[0.003]\n",
    "# epochs_vars = [100]#[50,100,200, 300, 500]#[50, 100, 200]\n",
    "# L1_lambda_vars = [1e-3]#[1e-4]#[1e-2,1e-3,1e-4,1e-5,1e-6]\n",
    "# L2_lambda_vars = [1e-2]#[1e-1,1e-2,1e-3,1e-4,1e-5]\n",
    "# weight_decay_vars = [9e-5]#[9e-2,9e-3,9e-4,9e-5,9e-6]#[9e-5] #5e-2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## CANNY\n",
    "# architecture_vars = ['a']\n",
    "# droput_vars = [.3]#[0,.4,.8]\n",
    "\n",
    "# steer_noise_level_vars = [10, 12,14] #[8, 12]#[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]#[0,4,8,10,12,16,20]#[0,2, 4,6, 8, 10, 12, 14, 16, 18, 20]\n",
    "# he_distance_vars = [.4,.5,.6] #[.5,.6]#[.3,.4,.5,.6,.7,.8,.9]#[.2,.3,.4,.5,.6,.7,.8,.9]#[.4,.6,.8]#[.3,.4,.5,.55,.6,.7,.8,.9]\n",
    "\n",
    "# canny1_vars = [100, 0]#[100, 0]\n",
    "# canny2_vars = [200, 0]#[200, 0]\n",
    "# blur_vars = [3]#[0,3,5,7,9]#[3]#[0,3,5,7,9]\n",
    "# img_noise_vars = [80]#[0,40,80,120,160]#[80]#[0,40,80,120,160]\n",
    "# keep_bottom_vars = [.8]#[.8]#[.4,.5,.6,.7,.8,.9]#[.8]#[.4,.5,.6,.7,.8,.9]\n",
    "# ds_length_vars = [10000]\n",
    "# img_size_vars = [32]\n",
    "\n",
    "# batch_size_vars = [2**16]#[2**8, 2**9,2**10,2**11,2**12,2**13,2**14,2**15,2**16]\n",
    "# lr_vars =  [5e-3]#[3e-3, 5e-3, 3e-3]#[1e-1, 5e-2, 1e-2, 5e-3, 3e-3, 1e-3, 1e-4, 1e-5]#[0.003]\n",
    "# epochs_vars = [100]#[50,100,200, 300, 500]#[50, 100, 200]\n",
    "# L1_lambda_vars = [1e-3]#[1e-4]#[1e-2,1e-3,1e-4,1e-5,1e-6]\n",
    "# L2_lambda_vars = [1e-2]#[1e-1,1e-2,1e-3,1e-4,1e-5]\n",
    "# weight_decay_vars = [9e-5]#[9e-2,9e-3,9e-4,9e-5,9e-6]#[9e-5] #5e-2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ### IMG NOISE\n",
    "# architecture_vars = ['a']\n",
    "# droput_vars = [.3]#[0,.4,.8]\n",
    "\n",
    "# steer_noise_level_vars = [10, 12] #[8, 12]#[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]#[0,4,8,10,12,16,20]#[0,2, 4,6, 8, 10, 12, 14, 16, 18, 20]\n",
    "# he_distance_vars = [.4,.5] #[.5,.6]#[.3,.4,.5,.6,.7,.8,.9]#[.2,.3,.4,.5,.6,.7,.8,.9]#[.4,.6,.8]#[.3,.4,.5,.55,.6,.7,.8,.9]\n",
    "\n",
    "# canny1_vars = [100]#[100, 0]\n",
    "# canny2_vars = [200]#[200, 0]\n",
    "# blur_vars = [3]#[0,3,5,7,9]#[3]#[0,3,5,7,9]\n",
    "# img_noise_vars = [0,40,80,120,160]#[80]#[0,40,80,120,160]\n",
    "# keep_bottom_vars = [.8]#[.8]#[.4,.5,.6,.7,.8,.9]#[.8]#[.4,.5,.6,.7,.8,.9]\n",
    "# ds_length_vars = [10000]\n",
    "# img_size_vars = [32]\n",
    "\n",
    "# batch_size_vars = [2**16]#[2**8, 2**9,2**10,2**11,2**12,2**13,2**14,2**15,2**16]\n",
    "# lr_vars =  [5e-3]#[3e-3, 5e-3, 3e-3]#[1e-1, 5e-2, 1e-2, 5e-3, 3e-3, 1e-3, 1e-4, 1e-5]#[0.003]\n",
    "# epochs_vars = [100]#[50,100,200, 300, 500]#[50, 100, 200]\n",
    "# L1_lambda_vars = [1e-3]#[1e-4]#[1e-2,1e-3,1e-4,1e-5,1e-6]\n",
    "# L2_lambda_vars = [1e-2]#[1e-1,1e-2,1e-3,1e-4,1e-5]\n",
    "# weight_decay_vars = [9e-5]#[9e-2,9e-3,9e-4,9e-5,9e-6]#[9e-5] #5e-2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## ARTICLE \n",
    "architecture_vars = ['a']\n",
    "droput_vars = [.3]\n",
    "\n",
    "steer_noise_level_vars =[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]\n",
    "he_distance_vars = [.2,.3,.4,.5,.6,.7,.8,.9]\n",
    "\n",
    "canny1_vars = [100]\n",
    "canny2_vars = [200]\n",
    "blur_vars = [3]\n",
    "img_noise_vars = [80]\n",
    "keep_bottom_vars = [.8]\n",
    "ds_length_vars = [10000]\n",
    "img_size_vars = [32]\n",
    "\n",
    "batch_size_vars = [2**16]\n",
    "lr_vars =  [5e-3,3e-3]\n",
    "epochs_vars = [100,150,200]\n",
    "L1_lambda_vars = [1e-3]\n",
    "L2_lambda_vars = [1e-2]\n",
    "weight_decay_vars = [9e-5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# he_distance_vars = [0.3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#generate datasets names\n",
    "#names and parameters\n",
    "datasets = []\n",
    "for steer_noise_level in steer_noise_level_vars:\n",
    "    for he_distance in he_distance_vars:\n",
    "        for canny1, canny2 in zip(canny1_vars, canny2_vars):\n",
    "            for blur in blur_vars:\n",
    "                for img_noise in img_noise_vars:\n",
    "                    for keep_bottom in keep_bottom_vars:\n",
    "                        for img_size in img_size_vars:\n",
    "                            for ds_length in ds_length_vars:\n",
    "                                name = f'ds_sn{steer_noise_level:.0f}_he{100*he_distance:.0f}_canny{canny1}_{canny2}_blur{blur:.0f}_noise{img_noise:.0f}_keep{100*keep_bottom:.0f}_size{img_size:.0f}_length{ds_length:.0f}'\n",
    "                                params = {'name':name, 'steer_noise_level': steer_noise_level, 'he_distance': he_distance, 'canny1': canny1, 'canny2': canny2, 'blur': blur, 'img_noise': img_noise, 'keep_bottom': keep_bottom, 'img_size': img_size, 'ds_length': ds_length}\n",
    "                                datasets.append(params)\n",
    "\n",
    "print(f'total dataset combinations: {len(datasets)}')\n",
    "\n",
    "all_names = []\n",
    "for ds in datasets:\n",
    "    all_names.append(ds['name'])\n",
    "#check if there are duplicates\n",
    "print(f'number of unique names: {len(set(all_names))}, number of names: {len(all_names)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PREPARE DATASETS\n",
    "for ds in tqdm(datasets):\n",
    "    prepare_ds(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # ANALYZE DATASETS\n",
    "# for ds in tqdm(datasets):\n",
    "#     analyze_ds(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#create starting network\n",
    "import torch.nn as nn\n",
    "\n",
    "epochs = 0\n",
    "max_imgs = 200000\n",
    "\n",
    "\n",
    "\n",
    "assert len(architecture_vars) == len(img_size_vars) == 1, 'only one architecture and one img_noise is allowed'\n",
    "for a,size in zip(architecture_vars, img_size_vars):\n",
    "    net_name = f'base_{a}_{size}'\n",
    "    if os.path.exists(f'tmp/models/{net_name}.pt'):\n",
    "        print('base network already exists')\n",
    "        break\n",
    "\n",
    "    all_imgs = []\n",
    "    all_hes = []\n",
    "    for ds_params in datasets:\n",
    "        ds_name = ds_params['name']\n",
    "        npz = my_load(f'tmp/dss/{ds_name}.npz', allow_pickle=True)\n",
    "        imgs, hes = npz['imgs'], npz['hes']\n",
    "        for img, he in zip(imgs, hes):\n",
    "            all_imgs.append(img)\n",
    "            all_hes.append(he)\n",
    "    all_imgs = np.array(all_imgs)\n",
    "    all_hes = np.array(all_hes)\n",
    "\n",
    "    #shuffle\n",
    "    idx = np.arange(len(all_imgs))\n",
    "    np.random.shuffle(idx)\n",
    "    all_imgs = all_imgs[idx]\n",
    "    all_hes = all_hes[idx]\n",
    "\n",
    "    #keep only max_imgs\n",
    "    all_imgs = all_imgs[:max_imgs]\n",
    "    all_hes = all_hes[:max_imgs]\n",
    "\n",
    "    big_ds_name = f'big_ds_{a}_{size}'\n",
    "    np.savez(f'tmp/dss/{big_ds_name}.npz', imgs=all_imgs, hes=all_hes, name=big_ds_name, img_size=size)\n",
    "    print(f'base network: {all_imgs.shape}, {all_hes.shape}')\n",
    "    ds = MyDataset(big_ds_name, device=device)\n",
    "        #create model\n",
    "    net = HEstimator()\n",
    "    net.apply(reset_weights)\n",
    "    net.to(device)\n",
    "\n",
    "    #create dataloader\n",
    "    train_size = int(0.8 * len(all_imgs))\n",
    "    val_size = len(ds) - train_size\n",
    "    train_ds, val_ds = torch.utils.data.random_split(ds, [train_size, val_size])\n",
    "\n",
    "    train_dataloader = DataLoader(train_ds, batch_size=65536, shuffle=True)\n",
    "    val_dataloader = DataLoader(val_ds, batch_size=65536, shuffle=False)\n",
    "\n",
    "    optimizer = torch.optim.Adam(net.parameters(), lr=0.003, weight_decay=9e-5)\n",
    "    regr_loss_fn1 = nn.MSELoss() #before epochs/2\n",
    "    regr_loss_fn2 = nn.MSELoss() #after epochs/2 for finetuning\n",
    "\n",
    "    #train\n",
    "\n",
    "    best_val = np.inf\n",
    "    best_epoch = 0\n",
    "    best_model = net\n",
    "    losses = np.zeros((epochs, 2))\n",
    "    for epoch in tqdm(range(epochs)):\n",
    "        regr_loss_fn = regr_loss_fn1 if epoch < epochs//2 else regr_loss_fn2\n",
    "        he_loss = train_epoch(net, train_dataloader, regr_loss_fn, optimizer, 1e-4, 1e-2)\n",
    "        val_he_loss = val_epoch(net, val_dataloader, regr_loss_fn)\n",
    "        losses[epoch, 0] = he_loss\n",
    "        losses[epoch, 1] = val_he_loss\n",
    "        clear_output(wait=False)\n",
    "        if val_he_loss < best_val:\n",
    "            best_val = val_he_loss\n",
    "            best_epoch = epoch\n",
    "            best_model = deepcopy(net)\n",
    "            print('saved model')\n",
    "            # torch.save(net.state_dict(), f'tmp/models/{net_name}.pt')\n",
    "        print(f'he_loss {he_loss:.4f}, val_he_loss {val_he_loss:.4f}, best_val {best_val:.4f}, best_epoch {best_epoch}')\n",
    "\n",
    "    #save losses\n",
    "    # np.save(f'tmp/{name}_losses.npy', losses)\n",
    "    torch.save(best_model.state_dict(), f'tmp/models/{net_name}.pt')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#create training combinations\n",
    "trainings_combinations = []\n",
    "for ds in datasets:\n",
    "    for architecture in architecture_vars:\n",
    "        for batch_size in batch_size_vars:\n",
    "            for lr in lr_vars:\n",
    "                for epochs in epochs_vars:\n",
    "                    for L1_lambda in L1_lambda_vars:\n",
    "                        for L2_lambda in L2_lambda_vars:\n",
    "                            for weight_decay in weight_decay_vars:\n",
    "                                for dropout in droput_vars:\n",
    "                                    name = f'tr_{ds[\"name\"]}_arch{architecture}_bs{batch_size:.0f}_lr{lr*10**6:.0f}_ep{epochs:.0f}_L1{L1_lambda*10**6:.0f}_L2{L2_lambda*10**6:.0f}_wd{weight_decay*10**6:.0f}_dr{dropout*100:.0f}'\n",
    "                                    params = {'name':name, 'ds_name': ds['name'], 'architecture': architecture, 'batch_size': batch_size, 'lr': lr, 'epochs': epochs, 'L1_lambda': L1_lambda, 'L2_lambda': L2_lambda, 'weight_decay': weight_decay, 'dropout': dropout}\n",
    "                                    trainings_combinations.append(params)\n",
    "\n",
    "all_names = []\n",
    "for tr in trainings_combinations:\n",
    "    all_names.append(tr['name'])\n",
    "#check if there are duplicates\n",
    "print(f'number of unique names: {len(set(all_names))}, number of names: {len(all_names)}')\n",
    "\n",
    "print(f'total training combinations: {len(trainings_combinations)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#check if all combinations are already there\n",
    "all_found = True\n",
    "not_founds = 0\n",
    "for tr in trainings_combinations:\n",
    "    tr_name = tr['name']\n",
    "    comb_path = f'tmp/training_combinations/{tr_name}.npz'\n",
    "    if not os.path.exists(comb_path):\n",
    "        # print(f'not found: {tr_name}')\n",
    "        all_found = False\n",
    "        not_founds += 1\n",
    "if all_found:\n",
    "    print('all combinations found')\n",
    "else:\n",
    "    print(f'{not_founds} combinations not found, time_estimate: {not_founds*20/60.0:.2f} mins')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#TRAINING \n",
    "for tr in tqdm(trainings_combinations):\n",
    "    train(tr, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# EVALUATION\n",
    "for tr in tqdm(trainings_combinations):\n",
    "    evaluate(tr, eval_datasets=ALL_EVALUATION_DATASETS, device=device, show_imgs=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Best MSEs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#BEST result REAL_NOISY_DATASETS+REAL_CLEAN_DATASETS\n",
    "best_combination, best_MSE, all_MSE = get_best_result(trainings_combinations, eval_datasets=ALL_REAL_EVAL_DATASETS, device=device)\n",
    "print(f'best combination: {best_combination}, best MSE: {best_MSE}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#BEST result SIM_DATASETS\n",
    "best_combination, best_MSE, all_MSE = get_best_result(trainings_combinations, eval_datasets=ALL_SIM_EVAL_DATASETS, device=device)\n",
    "print(f'best combination: {best_combination}, best MSE: {best_MSE}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for steer_noise_level \n",
    "steer_noise_level_MSEs = get_STDs_for('steer_noise_level', trainings_combinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for he_distance\n",
    "he_distance_MSEs = get_STDs_for('he_distance', trainings_combinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for lr\n",
    "lr_MSEs = get_STDs_for('lr', trainings_combinations, log=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for epochs\n",
    "epochs_MSEs = get_STDs_for('epochs', trainings_combinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for keep_bottom\n",
    "keep_bottom_MSEs = get_STDs_for('keep_bottom', trainings_combinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plto for img_noise\n",
    "img_noise_MSEs = get_STDs_for('img_noise', trainings_combinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for blur\n",
    "blur_MSEs = get_STDs_for('blur', trainings_combinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot for Canny\n",
    "Canny_MSEs = get_STDs_for('canny1', trainings_combinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for canny2\n",
    "Canny2_MSEs = get_STDs_for('canny2', trainings_combinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for weight_decay\n",
    "weight_decay_MSEs = get_STDs_for('weight_decay', trainings_combinations, log=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plor for L2_lambda\n",
    "L2_lambda_MSEs = get_STDs_for('L2_lambda', trainings_combinations, log=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for L1_lambda\n",
    "L1_lambda_MSEs = get_STDs_for('L1_lambda', trainings_combinations, log=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for dropout\n",
    "dropout_MSEs = get_STDs_for('dropout', trainings_combinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for batch_size\n",
    "batch_size_MSEs = get_STDs_for('batch_size', trainings_combinations, log=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Double Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for steer_noise_level and he_distance\n",
    "%matplotlib widget\n",
    "steer_noise_level_he_distance_MSEs = get2D_MSEs_for('steer_noise_level', 'he_distance', trainings_combinations, eval_datasets=REAL_NOISY_DATASETS+REAL_CLEAN_DATASETS, azimuth=-130, elevation=60, save=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for steer_noise_level and img_noise_vars\n",
    "%matplotlib widget\n",
    "steer_noise_level_img_noise_vars_MSEs = get2D_MSEs_for('steer_noise_level', 'img_noise', trainings_combinations, eval_datasets=REAL_NOISY_DATASETS+REAL_CLEAN_DATASETS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for he_distance and keep_bottom\n",
    "%matplotlib widget\n",
    "he_distance_keep_bottom_MSEs = get2D_MSEs_for('he_distance', 'keep_bottom', trainings_combinations, eval_datasets=REAL_NOISY_DATASETS+REAL_CLEAN_DATASETS, save=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for L2_lambda and weight_decay\n",
    "%matplotlib widget\n",
    "L2_lambda_weight_decay_MSEs = get2D_MSEs_for('L2_lambda', 'weight_decay', trainings_combinations, eval_datasets=REAL_NOISY_DATASETS+REAL_CLEAN_DATASETS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for lr and epochs\n",
    "%matplotlib widget\n",
    "lr_epochs_MSEs = get2D_MSEs_for('lr', 'epochs', trainings_combinations, eval_datasets=REAL_NOISY_DATASETS+REAL_CLEAN_DATASETS, azimuth=144, elevation=10, save=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name in all_names_used:\n",
    "    print(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CHECKS plot real and sim datasets\n",
    "\n",
    "# cv.namedWindow('image', cv.WINDOW_NORMAL)\n",
    "# cv.resizeWindow('image', 2*320, 240)\n",
    "he_distance_vars \n",
    "\n",
    "for real_name, sim_name in zip(REAL_NOISY_DATASETS+REAL_CLEAN_DATASETS, SIM_NOISY_DATASETS+SIM_CLEAN_DATASETS):\n",
    "    real = my_load(f'tmp/real_dss/{real_name}.npz')\n",
    "    sim = my_load(f'tmp/real_dss/{sim_name}.npz')\n",
    "\n",
    "    # rlocs = real['locs']\n",
    "    # slocs = sim['locs']\n",
    "\n",
    "    # rimgs = real['imgs']\n",
    "    # simgs = sim['imgs']\n",
    "\n",
    "    # #decimate\n",
    "    # rimgs = rimgs[::3]\n",
    "    # simgs = simgs[::3]\n",
    "\n",
    "    # for rimg, simg in zip(rimgs, simgs):\n",
    "    #     to_show = np.hstack((rimg, simg))\n",
    "\n",
    "    #     #add title in the imgs\n",
    "    #     cv.putText(to_show, real_name, (0, 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv.LINE_AA)\n",
    "    #     cv.imshow('image', to_show)\n",
    "    #     if cv.waitKey(1) == 27:\n",
    "    #         break\n",
    "    \n",
    "    for hed in he_distance_vars:\n",
    "        hed_name = f'{100*hed:.0f}'\n",
    "        real_hes = my_load(f'tmp/hes/{real_name}_{hed_name}.npz')['hes']\n",
    "        sim_hes = my_load(f'tmp/hes/{sim_name}_{hed_name}.npz')['hes']\n",
    "\n",
    "        assert np.allclose(real_hes, sim_hes) #they should be the same\n",
    "\n",
    "\n",
    "\n",
    "    # fig, ax = plt.subplots( figsize=(10,5))\n",
    "    # ax.plot(rlocs[:, 1], rlocs[:, 0], label='real', linewidth=1)\n",
    "    # ax.plot(slocs[:, 1], slocs[:, 0], label='sim', linewidth=1)\n",
    "    # ax.legend()\n",
    "    # ax.set_title(real_name)\n",
    "    # ax.set_aspect('equal')\n",
    "    # plt.show()\n",
    "\n",
    "# cv.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raise SystemExit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !!! CLEARING DATA !!!\n",
    "import shutil\n",
    "import cv2 as cv\n",
    "import os\n",
    "folders = ['dss', 'evals', 'hes', 'models', 'real_dss', 'to_del', 'training_combinations']\n",
    "folders = ['to_del']\n",
    "\n",
    "cv.namedWindow('CAREFUL', cv.WINDOW_NORMAL)\n",
    "cv.imshow('CAREFUL', np.zeros((100, 100, 3), dtype=np.uint8))\n",
    "cv.waitKey(0)\n",
    "cv.destroyAllWindows()\n",
    "\n",
    "cv.namedWindow('ARE U SURE U WANT TO DELETE EVERYTHING?', cv.WINDOW_NORMAL)\n",
    "cv.imshow('ARE U SURE U WANT TO DELETE EVERYTHING?', np.zeros((100, 100, 3), dtype=np.uint8))\n",
    "cv.waitKey(0)\n",
    "cv.destroyAllWindows()\n",
    "\n",
    "\n",
    "for folder in folders:\n",
    "    folder_path = f'tmp/{folder}'\n",
    "    shutil.rmtree(folder_path)\n",
    "    os.mkdir(folder_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raise SystemExit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# testN = HEstimator()\n",
    "# a = np.array([1,2,3,4,5,6,7,8,9,10])\n",
    "# print(testN)\n",
    "# print(a)\n",
    "# np.savez('testN', testN=testN, a=a)\n",
    "import numpy as np\n",
    "loadedN = np.load('testN.npz', allow_pickle=True)['testN']\n",
    "loadeda = np.load('testN.npz', allow_pickle=True)['a']\n",
    "print(loadedN)\n",
    "print(loadeda)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lane_keeper_ahead = HEstimator()\n",
    "\n",
    "#load pretrained model\n",
    "net_name = 'tmp_bk/models/tr_ds_sn12_he60_canny100_200_blur3_noise80_keep70_size32_length10000_archa_bs65536_lr3000_ep100_L11_L210000_wd90_dr30.pt'\n",
    "lane_keeper_ahead.load_state_dict(torch.load(net_name))\n",
    "\n",
    "lane_keeper_ahead.to(device)\n",
    "\n",
    "# name_dataset = 'big_ds_a_80' #'saved_tests/train18' #'saved_tests/sim_dataset0'\n",
    "#create dataset #takes a long time but then training is faster\n",
    "name_dataset = 'ds_sn14_he60_canny100_200_blur3_noise80_keep80_size32_length10000'\n",
    "train_dataset = MyDataset(name_dataset, device=device)\n",
    "\n",
    "#split dataset into train and val\n",
    "train_size = int(0.9*len(train_dataset))\n",
    "val_size = len(train_dataset) - train_size\n",
    "train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])\n",
    "\n",
    "# DATALOADERS\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=8*8192, shuffle=True)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=8192, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TRAINING \n",
    "#parameters\n",
    "lr = 0.005 #0.005\n",
    "epochs = 300 #500\n",
    "#regularization is applied only to convolutional section, add weight decay to apply it to all layers\n",
    "L1_lambda = 1e-4 #9e-4\n",
    "L2_lambda = 2e-2 #1e-2\n",
    "optimizer = torch.optim.Adam(lane_keeper_ahead.parameters(), lr=lr, weight_decay=9e-5) #wd = 2e-3# 3e-5\n",
    "regr_loss_fn1 = nn.MSELoss() #before epochs/2\n",
    "regr_loss_fn2 = nn.MSELoss() #after epochs/2 for finetuning\n",
    "\n",
    "best_val = 100\n",
    "best_epoch = 0\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    # try:\n",
    "    if True:\n",
    "        regr_loss_fn = regr_loss_fn1 if epoch < epochs//2 else regr_loss_fn2\n",
    "        he_loss = train_epoch(lane_keeper_ahead, train_dataloader, regr_loss_fn, optimizer, L1_lambda, L2_lambda)\n",
    "        val_he_loss = val_epoch(lane_keeper_ahead, val_dataloader, regr_loss_fn)\n",
    "        clear_output(wait=False)\n",
    "    # except Exception as e:\n",
    "    #     print(e)\n",
    "    #     torch.cuda.empty_cache()\n",
    "    #     continue\n",
    "    if val_he_loss < best_val:\n",
    "        best_val = val_he_loss\n",
    "        best_epoch = epoch\n",
    "        torch.save(lane_keeper_ahead.state_dict(), model_name)\n",
    "        print(\"model saved\")\n",
    "    \n",
    "    print(f\"Epoch  {epoch+1}/{epochs},  loss = {regr_loss_fn} \\nhe_loss: {he_loss:.4f},   Val: {val_he_loss:.4f}, best_val: {best_val:.4f}, best_epoch: {best_epoch}\")\n",
    "    # print(f\"lat_err_loss2: {err_loss2:.4f},   Val: {val_loss2:.4f}\")\n",
    "    # print(f\"curv_loss: {curv_loss}\")\n",
    "\n",
    "#Note: sweet spot for training is around 0.016 -> 0.020, also note that training can get stuck, and loss can start improving randomly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# EVALUATE ON TEST SET (UNSEEN DATA)\n",
    "lane_keeper_ahead.load_state_dict(torch.load(model_name))\n",
    "he_loss = val_epoch(lane_keeper_ahead, val_dataloader, regr_loss_fn)\n",
    "\n",
    "# print(f\"lateral_err2_loss: {err_loss2}\")\n",
    "print(f\"he loss: {he_loss}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# VISUALIZE CONVOLUTIONAL FILTERS\n",
    "conv_layers = []\n",
    "children = list(lane_keeper_ahead.children())\n",
    "for i in range(len(children)):\n",
    "    if isinstance(children[i], nn.Conv2d):\n",
    "        conv_layers.append(children[i])\n",
    "    elif isinstance(children[i], nn.Sequential):\n",
    "        for child in children[i].children():\n",
    "            if isinstance(child, nn.Conv2d):\n",
    "                conv_layers.append(child)\n",
    "\n",
    "c0 = conv_layers[0].weight.data.cpu().numpy()\n",
    "c1 = conv_layers[1].weight.data.cpu().numpy()\n",
    "c2 = conv_layers[2].weight.data.cpu().numpy()\n",
    "\n",
    "def plot_nchw_data(data, h_num, v_num, title, size=(10, 10)):\n",
    "    fig, axs = plt.subplots(h_num, v_num, figsize=size)\n",
    "    shape = data.shape\n",
    "    data = data.reshape(shape[0]*shape[1], shape[2], shape[3])\n",
    "    for idx, ax in enumerate(axs.flatten()):\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])\n",
    "        if idx < len(data):\n",
    "            ax.imshow(data[idx,:,:], cmap='gray')\n",
    "    plt.suptitle(title)\n",
    "    #plt.tight_layout(rect=[0, 0, 1, 0.97], h_pad=0, w_pad=0)\n",
    "    plt.show()\n",
    "    return fig\n",
    "\n",
    "# fig0 = plot_nchw_data(c0, 4, 4, 'conv0')\n",
    "print(c0.shape)\n",
    "print(c1.shape)\n",
    "print(c2.shape)\n",
    "\n",
    "fig0 = plot_nchw_data(c0, 1, 4, 'conv0', size=(8,2))\n",
    "\n",
    "fig1 = plot_nchw_data(c1, 4, 4, 'conv1', size=(5,5)) \n",
    "\n",
    "fig2 = plot_nchw_data(c2, 8, 8, 'conv2', size=(10,10))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CONVERT TO ONNX MODEL FOR OPENCV\n",
    "lane_keeper_ahead.load_state_dict(torch.load(model_name))\n",
    "\n",
    "#save the model so that opencv can load it\n",
    "import torch\n",
    "import torch.onnx\n",
    "import torchvision\n",
    "import torchvision.models as models\n",
    "\n",
    "device = torch.device('cpu')\n",
    "lane_keeper_ahead.to(device)\n",
    "\n",
    "# set the model to inference mode\n",
    "lane_keeper_ahead.eval()\n",
    "\n",
    "# Create some sample input in the shape this model expects \n",
    "# This is needed because the convertion forward pass the network once \n",
    "dummy_input = torch.randn(1, 1, 32, 32)\n",
    "torch.onnx.export(lane_keeper_ahead, dummy_input, onnx_lane_keeper_path, verbose=True)\n",
    "\n",
    "clear_output(wait=False)\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "lane_keeper_ahead.to(device)"
   ]
  }
 ],
 "metadata": {
  "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.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}