train_ffoom_mlp.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "24c35a1e",
   "metadata": {},
   "source": [
    "# Train MLP on the FFOOM dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f589830f-9bf9-4a1b-8d09-a735b8d6d0b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: XLA_PYTHON_CLIENT_PREALLOCATE=false\n"
     ]
    }
   ],
   "source": [
    "%load_ext lab_black\n",
    "%env XLA_PYTHON_CLIENT_PREALLOCATE=false\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import pickle\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from typing import List, Tuple\n",
    "import seaborn as sns\n",
    "import sys\n",
    "import wandb\n",
    "\n",
    "import brunoflow as bf\n",
    "from brunoflow.ad import Node\n",
    "from brunoflow.net import MLP\n",
    "from brunoflow.opt import Adam, cross_entropy_loss, regularize\n",
    "from preprocessing.datasets import MNIST, FFOOM\n",
    "from utils import catchtime, gpu_memory_usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1d08b9c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running JAX on NVIDIA GeForce RTX 3050 Ti Laptop GPU\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    import jax\n",
    "    from jax import numpy as jnp\n",
    "\n",
    "    jax_device_kind = jax.devices()[0].device_kind\n",
    "    print(f\"Running JAX on {jax_device_kind}\")\n",
    "    # if \"NVIDIA\" not in jax_device_kind:\n",
    "    #     raise ValueError(\"Imported JAX, but not running on a GPU, terminating.\")\n",
    "except ImportError:\n",
    "    print(\"No JAX available to import!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "53cb83df",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"WANDB_NOTEBOOK_NAME\"] = os.path.join(os.getcwd(), \"train_ffoom_mlp.ipynb\")\n",
    "excluded_all_caps_params = {k for k in locals().keys() if k.isupper()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2a63e06f-5fd8-46bf-8135-8ee47f219348",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "# Data parameters\n",
    "DATASET_NAME = \"FFOOM\"\n",
    "DATASET_KWARGS_IDENTIFIABLE = {\"num_points\": 10000}\n",
    "DATASET_KWARGS = {}\n",
    "SEED = 2221\n",
    "BATCH_SIZE = 32\n",
    "TEST_BATCH_SIZE = 100\n",
    "SAMPLING_FRACTION = 1.0\n",
    "\n",
    "# Model parameters\n",
    "INPUT_SIZE = 4  # 28x28\n",
    "HIDDEN_SIZE = 8\n",
    "NUM_CLASSES = 2\n",
    "NUM_EPOCHS = 1\n",
    "LEARNING_RATE = 0.001\n",
    "L1_WEIGHT = 0.0\n",
    "L2_WEIGHT = 0.0\n",
    "DROPOUT_PROB = 0.0\n",
    "OVERWRITE_MODEL = True\n",
    "\n",
    "# Analysis parameters\n",
    "NUM_ELEMENTS_PER_STRATA = 32\n",
    "\n",
    "# Run parameters\n",
    "PM_RUN_ID = \"run_id\"\n",
    "PROJECT_NAME = \"bauer-ffoom\"\n",
    "# GROUP_NAME = \"regularized-loss\"\n",
    "# GROUP_NAME = \"scratch\"\n",
    "# GROUP_NAME = \"visualization\"\n",
    "GROUP_NAME = \"dropout\"\n",
    "TAGS = [\"jax\", \"dropout\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6a10f46d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct dataset and data ids\n",
    "dataset = getattr(sys.modules[__name__], DATASET_NAME)(\n",
    "    **{**DATASET_KWARGS_IDENTIFIABLE, **DATASET_KWARGS}\n",
    ")\n",
    "data_id = f\"{dataset.get_name()}\"\n",
    "data_dir = os.path.join(\"data\", DATASET_NAME, data_id, f\"{SAMPLING_FRACTION}-{SEED}\")\n",
    "input_dir = os.path.join(data_dir, \"inputs\")\n",
    "train_data_path = os.path.join(input_dir, \"train.pt\")\n",
    "val_data_path = os.path.join(input_dir, \"val.pt\")\n",
    "test_data_path = os.path.join(input_dir, \"test.pt\")\n",
    "\n",
    "# Construct model id\n",
    "model_id = f\"MLP-hidden_sz{HIDDEN_SIZE}-bs{BATCH_SIZE}-lr{LEARNING_RATE}-n{NUM_EPOCHS}\"\n",
    "model_id += f\"-l1_weight{L1_WEIGHT}\" if L1_WEIGHT != 0 else \"\"\n",
    "model_id += f\"-l2_weight{L2_WEIGHT}\" if L2_WEIGHT != 0 else \"\"\n",
    "model_id += f\"-dropoutprob{DROPOUT_PROB}\" if DROPOUT_PROB != 0 else \"\"\n",
    "\n",
    "model_dir = os.path.join(data_dir, \"models\", model_id)\n",
    "\n",
    "# Set random seeds\n",
    "torch.manual_seed(SEED)\n",
    "np.random.seed(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "94b8d015",
   "metadata": {},
   "outputs": [],
   "source": [
    "params_to_log = {\n",
    "    k: v\n",
    "    for k, v in locals().items()\n",
    "    if k.isupper() and k not in excluded_all_caps_params\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f9f20d93",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mkdu\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
      "wandb: ERROR Failed to sample metric: Not Supported\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.13.9 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.13.5"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/kevin/code/rycolab/interpreta-bauer-ly/wandb/run-20230203_192813-2u32blit</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/kdu/bauer-ffoom/runs/2u32blit\" target=\"_blank\">glistening-monkey-628</a></strong> to <a href=\"https://wandb.ai/kdu/bauer-ffoom\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'DATASET_NAME': 'FFOOM', 'DATASET_KWARGS_IDENTIFIABLE': {'num_points': 10000}, 'DATASET_KWARGS': {}, 'SEED': 2221, 'BATCH_SIZE': 32, 'TEST_BATCH_SIZE': 100, 'SAMPLING_FRACTION': 1.0, 'INPUT_SIZE': 4, 'HIDDEN_SIZE': 8, 'NUM_CLASSES': 2, 'NUM_EPOCHS': 1, 'LEARNING_RATE': 0.001, 'L1_WEIGHT': 0.0, 'L2_WEIGHT': 0.0, 'DROPOUT_PROB': 0.0, 'OVERWRITE_MODEL': True, 'NUM_ELEMENTS_PER_STRATA': 32, 'PM_RUN_ID': 'run_id', 'PROJECT_NAME': 'bauer-ffoom', 'GROUP_NAME': 'dropout', 'TAGS': ['jax', 'dropout']}\n"
     ]
    }
   ],
   "source": [
    "run = wandb.init(\n",
    "    project=PROJECT_NAME,\n",
    "    group=GROUP_NAME,\n",
    "    config=params_to_log,\n",
    "    tags=TAGS,\n",
    ")\n",
    "print(dict(wandb.config))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5914462",
   "metadata": {},
   "source": [
    "### Data Retrieval and Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3c23e381",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading cached train and test sets from data/FFOOM/FFOOM-num_points10000/1.0-2221/inputs/train.pt and data/FFOOM/FFOOM-num_points10000/1.0-2221/inputs/test.pt.\n"
     ]
    }
   ],
   "source": [
    "if (\n",
    "    os.path.isfile(train_data_path)\n",
    "    and os.path.isfile(val_data_path)\n",
    "    and os.path.isfile(test_data_path)\n",
    "):\n",
    "    print(\n",
    "        f\"Loading cached train and test sets from {train_data_path} and {test_data_path}.\"\n",
    "    )\n",
    "    train_data = torch.load(train_data_path)\n",
    "    val_data = torch.load(val_data_path)\n",
    "    test_data = torch.load(test_data_path)\n",
    "else:\n",
    "    train_data = dataset.get_train_data()\n",
    "    val_data = dataset.get_val_data()\n",
    "    test_data = dataset.get_test_data()\n",
    "\n",
    "    # Save the dataset\n",
    "    os.makedirs(input_dir, exist_ok=True)\n",
    "    torch.save(train_data, train_data_path)\n",
    "    torch.save(val_data, val_data_path)\n",
    "    torch.save(test_data, test_data_path)\n",
    "\n",
    "train_kwargs = {\"batch_size\": BATCH_SIZE}\n",
    "val_kwargs = {\"batch_size\": len(val_data)}\n",
    "test_kwargs = {\"batch_size\": TEST_BATCH_SIZE}\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_data, **train_kwargs)\n",
    "val_loader = torch.utils.data.DataLoader(val_data, **val_kwargs)\n",
    "test_loader = torch.utils.data.DataLoader(test_data, **test_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7faaff62",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Adding directory to artifact (./data/FFOOM/FFOOM-num_points10000/1.0-2221/inputs)... Done. 0.0s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logging datasets to w&b run <wandb.sdk.wandb_run.Run object at 0x7fa4132f5c70>.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<wandb.sdk.wandb_artifacts.Artifact at 0x7fa4132dd2b0>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# After loading/preprocessing your dataset, log it as an artifact to W&B\n",
    "print(f\"Logging datasets to w&b run {wandb.run}.\")\n",
    "artifact = wandb.Artifact(name=data_id, type=\"dataset\")\n",
    "artifact.add_dir(local_path=input_dir)\n",
    "run.log_artifact(artifact)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "df94ac6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAi8AAAEfCAYAAACJRiZ0AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8o6BhiAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAeOUlEQVR4nO3df2xV9f3H8Vcp9LZMehuo/YGUUrOqXbsZuSgUrUZNrkBkgi7DmdQ5AdMNNNCQTGQJSobNlsUQv1IMFYgIji5BJipBmmErrq0/mmawpDYmU1tKy+/eAtNbac/3D9aOw739Ne/pPZ/b5yM5f9xzP7f3fW9fbV/cezknzrIsSwAAAIYYF+0BAAAARoLyAgAAjEJ5AQAARqG8AAAAo1BeAACAUSgvAADAKJQXAABgFMoLAAAwCuUFAAAYhfICAACM4mh5OX/+vIqLi+X1euX1elVcXKzOzs5Bb/PEE08oLi7Ots2ZM8fJMYEQZBemIrsYC8Y7+cUfe+wxHT9+XAcPHpQkPfXUUyouLtY777wz6O3mzZunHTt29F9OSEhwckwgBNmFqcguxgLHyktTU5MOHjyo+vp6zZ49W5JUUVGhwsJCNTc36+abbx7wth6PRxkZGU6NBgyK7MJUZBdjhWPlpa6uTl6vt/8HSJLmzJkjr9er2traQX+IqqurlZaWppSUFN1zzz3auHGj0tLSwq4NBoMKBoP9l3t7e3Xu3DlNmTJFcXFxkXtAGDMOHz6s5ORk3XDDDert7dW4ceMcya5EfhFZZBcmsyxLFy5c0NSpUzVu3BCfarEcsnHjRis3Nzdkf25urvXiiy8OeLs9e/ZY7777rnXs2DFr//791q233mrl5+db3377bdj169evtySxsTmytba2OpZd8svm5EZ22Uzdrs7uQOIsy7I0As8//7xeeOGFQdd8+umnOnTokF5//XU1NzfbrsvNzdXSpUv17LPPDuv+2tvblZ2drT179ujhhx8Ouf7a9h8IBDR9+nQtXrxYEyZMGNZ9OO0vf/lLtEcIEQgEoj2Czccff+z4fezcuVO7du0adM0rr7yihoYGvf/++zpx4oQ6Ozvl9XolRT670sD5bW1tVXJy8jAfmbP6Hr+b/PGPf4z2CDaNjY2Ofv1jx47pn//856Br/H6/Ojo69K9//UsXL16MWnYrKio0ceLEYT4yZ/3mN7+J9ggh3Pa7103zdHV1KSsry5bdgYz4baOVK1fq0UcfHXTNjBkzdPToUZ08eTLkutOnTys9PX3Y95eZmans7Gx98cUXYa/3eDzyeDwh+ydMmOCa8uJGbvnD2OcHP/iB4/fxi1/8QgsWLBh0TUZGhtra2vp/oK9++TvS2ZUGzm9ycrLrvkdukpSUFO0RbJz+cGt+fr5++MMfDrrmuuuu08WLF/sLRbSyO3HiRNeUF96+Gpobf88M5/s24vKSmpqq1NTUIdcVFhYqEAjok08+0R133CHpyr+uA4GA5s6dO+z7O3v2rFpbW5WZmTnSUQGblJQUpaSkDLmuoKBAly5dsu0ju4imxMREJSYmDrnu+uuv13fffWfbR3YRixw7zkteXp7mzZun5cuXq76+XvX19Vq+fLkefPBB24fGbrnlFu3bt0+SdPHiRa1Zs0Z1dXX66quvVF1drYULFyo1NVWLFy92alTAZsaMGZo1a5akK2+Bkl2YIiUlpb9wkF3EMkcPUrd79279+Mc/lt/vl9/v109+8hO98cYbtjXNzc39L9HHx8fr2LFjeuihh3TTTTfpl7/8pW666SbV1dVp0qRJTo4K2PR9NmDx4sVkF0YpLCyURHYR2xw9SN3kyZOH/IDk1Z8XTkpK0vvvv+/kSMCw9L0PfPz48QHfEya7cKO+z6GQXcQyzm0EAACMQnkBAABGobwAAACjUF4AAIBRKC8AAMAolBcAAGAUygsAADDKqJSX8vJy5eTkKDExUT6fT0eOHBl0fU1NjXw+nxITE3XjjTfq1VdfHY0xgRAVFRVkF0Yiu4hljpeXyspKrVq1SuvWrVNjY6OKioo0f/58tbS0hF3/5ZdfasGCBSoqKlJjY6Oee+45PfPMM9q7d6/TowIh1q5dS3ZhJLKLWOZ4eXnppZe0dOlSLVu2THl5edq0aZOysrK0ZcuWsOtfffVVTZ8+XZs2bVJeXp6WLVumJ598Un/605+cHhUIUVxcTHZhJLKLWOZoeenu7lZDQ4P8fr9tv9/vV21tbdjb1NXVhax/4IEH9Nlnn4WcLVWSgsGgurq6bBvwffVl7b777rPtj2R2JfKLyOvp6ZFEdhHbHC0vZ86cUU9Pj9LT023709PT1dHREfY2HR0dYddfvnxZZ86cCVlfVlYmr9fbv2VlZUXuAWDM6vtFnJaWZtsfyexK5BeRFwwGJZFdxLZR+cBuXFyc7bJlWSH7hlofbr905X3dQCDQv7W2tkZgYuAKJ7MrkV84h+wiljl6VunU1FTFx8eHtP1Tp06FtPw+GRkZYdePHz9eU6ZMCVnv8Xj6z6IKRErf2XhPnjxp2x/J7ErkF5HXlyeyi1jm6CsvCQkJ8vl8qqqqsu2vqqrS3Llzw96msLAwZP2hQ4c0a9YsTZgwwbFZgav1Ze2DDz6w7Se7cLv4+HhJZBexzfG3jUpLS/Xaa69p+/btampq0urVq9XS0qKSkhJJV156fPzxx/vXl5SU6Ouvv1Zpaamampq0fft2bdu2TWvWrHF6VCDEzp07yS6MRHYRyxx920iSlixZorNnz2rDhg1qb29XQUGBDhw4oOzsbElSe3u77dgDOTk5OnDggFavXq3Nmzdr6tSpevnll/XII484PSoQoqysjOzCSGQXsSzO6vtUVozo6uqS1+vVz3/+c9e83Ll79+5ojxDCbd/2gf4LZ7RcunRJfr9fgUCg//Mvo6Evv6N9v4MZ7EOe0fJ///d/0R7B5rPPPov2CP26u7v15z//OWrZ3b17tyZOnDhq9zuYX/3qV9EeIURnZ2e0R7Bx09+Ckfz+49xGAADAKJQXAABgFMoLAAAwCuUFAAAYhfICAACMQnkBAABGobwAAACjjEp5KS8vV05OjhITE+Xz+XTkyJEB11ZXVysuLi5k+/zzz0djVMCmoqKC7MJIZBexzPHyUllZqVWrVmndunVqbGxUUVGR5s+fbzu6YzjNzc1qb2/v33Jzc50eFQixdu1asgsjkV3EMsfLy0svvaSlS5dq2bJlysvL06ZNm5SVlaUtW7YMeru0tDRlZGT0b30nGwNGU3FxMdmFkcguYpmj5aW7u1sNDQ3y+/22/X6/f8jDwd92223KzMzU/fffH3J2VMBp3333nSTpvvvus+0nu3C7np4eSWQXsc3REzOeOXNGPT09Sk9Pt+1PT09XR0dH2NtkZmZq69at8vl8CgaDeuONN3T//ferurpad999d8j6YDCoYDDYfzkQCEj67x8fhNfV1RXtEWwuXboU7RFs+vJ5/fXX2/ZHMrvSwPl12/fHbb755ptoj2DT3d0d7RH69f0sRSu7//73vyPxMCLCTeftcSs3/a7pm2VY3zfLQW1tbZYkq7a21rb/97//vXXzzTcP++s8+OCD1sKFC8Net379eksSG5sj21//+lfHskt+2ZzcyC6bqVtra+uQ+XT0lZfU1FTFx8eHtP1Tp06FvBozmDlz5mjXrl1hr1u7dq1KS0v7L/f29urcuXOaMmXK9z4bbldXl7KystTa2uqaM/y6Taw+R8FgUJmZmbp8+bJtfySzKzmX31j9vkRSrD5HZDf2xepzZFmWLly4oKlTpw651tHykpCQIJ/Pp6qqKi1evLh/f1VVlR566KFhf53GxkZlZmaGvc7j8cjj8dj2paSk/E/zDiQ5OTmmAuKEWHyOfD6f/va3v+mRRx7p3xfJ7ErO5zcWvy+RFovPEdkdG2LxOfJ6vcNa52h5kaTS0lIVFxdr1qxZKiws1NatW9XS0qKSkhJJV9p7W1ubdu7cKUnatGmTZsyYofz8fHV3d2vXrl3au3ev9u7d6/SogA3ZhanILmKd4+VlyZIlOnv2rDZs2KD29nYVFBTowIEDys7OliS1t7fbjj3Q3d2tNWvWqK2tTUlJScrPz9d7772nBQsWOD0qYEN2YSqyi1gXZ1l8HHsgwWBQZWVlWrt2bcjLo7iC58id+L4MjefInfi+DI3niPICAAAMw4kZAQCAUSgvAADAKJQXAABgFMoLAAAwCuVlAOXl5crJyVFiYqJ8Pp+OHDkS7ZFco6ysTLfffrsmTZqktLQ0LVq0SM3NzdEeC/9BdgdGdt2N7A6M7NpRXsKorKzUqlWrtG7dOjU2NqqoqEjz58+3HRdhLKupqdGKFStUX1+vqqoqXb58WX6/33UnVxyLyO7gyK57kd3BkV07/qt0GLNnz9bMmTO1ZcuW/n15eXlatGiRysrKojiZO50+fVppaWmqqakZ8Ay0GB1kd2TIrnuQ3ZEZ69nllZdrdHd3q6GhQX6/37bf7/ertrY2SlO5WyAQkCRNnjw5ypOMbWR35MiuO5DdkRvr2aW8XOPMmTPq6ekJOftqenp6yNmxceUsoKWlpbrrrrtUUFAQ7XHGNLI7MmTXPcjuyJDdUTi3kamuPaW7ZVnf6zTvsWrlypU6evSoPvroo2iPgv8gu8NDdt2H7A4P2aW8hEhNTVV8fHxI2z916lTIvwrGuqefflr79+/Xhx9+qGnTpkV7nDGP7A4f2XUXsjt8ZPcK3ja6RkJCgnw+n6qqqmz7q6qqNHfu3ChN5S6WZWnlypV66623dPjwYeXk5ER7JIjsDgfZdSeyOzSya8crL2GUlpaquLhYs2bNUmFhobZu3aqWlhaVlJREezRXWLFihd588029/fbbmjRpUv+/lrxer5KSkqI83dhGdgdHdt2L7A6O7F7DQlibN2+2srOzrYSEBGvmzJlWTU1NtEdyDUlhtx07dkR7NFhkdzBk193I7sDIrh3HeQEAAEbhMy8AAMAolBcAAGAUygsAADAK5QUAABiF8gIAAIxCeQEAAEahvAAAAKNQXgAAgFEoLwAAwCiUFwAAYBTKCwAAMArlBQAAGIXyAgAAjEJ5AQAARqG8AAAAo1BeAACAUSgvAADAKJQXAABgFMoLAAAwCuUFAAAYhfICAACMQnkBAABGobwAAACjUF4AAIBRKC8AAMAolBcAAGAUygsAADAK5QUAABiF8gIAAIxCeQEAAEahvAAAAKNQXgAAgFEoLwAAwCiUFwAAYBTKCwAAMArlBQAAGIXyAgAAjEJ5AQAARqG8AAAAo1BeAACAUSgvAADAKJQXAABgFMoLAAAwCuUFAAAYhfICAACMQnkBAABGobwAAACjUF4AAIBRKC8AAMAolBcAAGAUygsAADAK5QUAABiF8gIAAIxCeQEAAEahvAAAAKNQXgAAgFEoLwAAwCiUFwAAYBTKCwAAMArlBQAAGIXyAgAAjEJ5AQAARqG8AAAAo1BeAACAUSgvAADAKJQXAABgFMoLAAAwCuUFAAAYhfICAACMQnkBAABGobwAAACjUF4AAIBRKC8AAMAolBcAAGAUygsAADAK5QUAABiF8gIAAIxCeQEAAEahvAAAAKNQXgAAgFEoLwAAwCiUFwAAYBTKCwAAMArlBQAAGIXyAgAAjEJ5AQAARqG8AAAAo1BeAACAUSgvAADAKJQXAABgFMoLAAAwCuUFAAAYhfICAACMQnkBAABGobwAAACjUF4AAIBRKC8AAMAolBcAAGAUygsAADAK5QUAABiF8gIAAIxCeQEAAEahvAAAAKNQXgAAgFEoLwAAwCiUFwAAYBTKCwAAMArlBQAAGIXyAgAAjEJ5AQAARqG8AAAAo1BeAACAUSgvAADAKJQXAABgFMoLAAAwCuUFAAAYhfICAACMQnkBAABGobwAAACjUF4AAIBRKC8AAMAolBcAAGAUygsAADAK5QUAABiF8gIAAIxCeQEAAEahvAAAAKNQXgAAgFEoLwAAwCiUFwAAYBTKCwAAMArlBQAAGIXyAgAAjEJ5AQAARqG8AAAAo1BeAACAUSgvAADAKJQXAABgFMoLAAAwiqPl5fz58youLpbX65XX61VxcbE6OzsHvc0TTzyhuLg42zZnzhwnxwRCkF2YiuxiLBjv5Bd/7LHHdPz4cR08eFCS9NRTT6m4uFjvvPPOoLebN2+eduzY0X85ISHByTGBEGQXpiK7GAscKy9NTU06ePCg6uvrNXv2bElSRUWFCgsL1dzcrJtvvnnA23o8HmVkZDg1GjAosgtTkV2MFY6Vl7q6Onm93v4fIEmaM2eOvF6vamtrB/0hqq6uVlpamlJSUnTPPfdo48aNSktLC7s2GAwqGAz2X+7t7dW5c+c0ZcoUxcXFRe4BYcw4fPiwkpOTdcMNN6i3t1fjxo1zJLsS+UVkkV2YzLIsXbhwQVOnTtW4cUN8qsVyyMaNG63c3NyQ/bm5udaLL7444O327Nljvfvuu9axY8es/fv3W7feequVn59vffvtt2HXr1+/3pLExubI1tra6lh2yS+bkxvZZTN1uzq7A4mzLMvSCDz//PN64YUXBl3z6aef6tChQ3r99dfV3Nxsuy43N1dLly7Vs88+O6z7a29vV3Z2tvbs2aOHH3445Ppr238gEND06dOVlpY2dHMbJdc+B27Q09MT7RFsJk+eHO0Rwurs7JTX65UU+exKA+e3tbVVycnJ3/8BREDf43eT2traaI9g4/RztHnzZpWXlw+6prKyUrW1tdq3b59aWlqilt1XXnlFSUlJw3xkzvrZz34W7RFCuO3nadq0adEeoV9vb69OnDhhy+5ARvy20cqVK/Xoo48OumbGjBk6evSoTp48GXLd6dOnlZ6ePuz7y8zMVHZ2tr744ouw13s8Hnk8npD948aNc015ccsfoau5rby41dUvf0c6u9LA+U1OTnZlbtziuuuui/YINpMmTXL06z/11FND/iGeNm2aWlpadP78eUnRy25SUpImTpw47PtyEj9DQ3PL38mrDedtxxGXl9TUVKWmpg65rrCwUIFAQJ988onuuOMOSdLHH3+sQCCguXPnDvv+zp49q9bWVmVmZo50VCBiyC6iafLkycN6dXLmzJm6cOGCbR/ZRSxyrHLl5eVp3rx5Wr58uerr61VfX6/ly5frwQcftH1o7JZbbtG+ffskSRcvXtSaNWtUV1enr776StXV1Vq4cKFSU1O1ePFip0YFBvTpp5+SXRgjNzdXd911lySyi9jm6HFedu/erWeeeUZ+v1+S9NOf/lSvvPKKbU1zc7MCgYAkKT4+XseOHdPOnTvV2dmpzMxM3XvvvaqsrHT8ZVkgnL5f3mQXpvjDH/6gO++8k+wipjlaXiZPnqxdu3YNuubqzwsnJSXp/fffd3IkYESOHz8+4PvmZBdulJKSIonsIra575M6AAAAg6C8AAAAo1BeAACAUSgvAADAKJQXAABgFMoLAAAwCuUFAAAYZVTKS3l5uXJycpSYmCifz6cjR44Mur6mpkY+n0+JiYm68cYb9eqrr47GmECIiooKsgsjkV3EMsfLS2VlpVatWqV169apsbFRRUVFmj9/vlpaWsKu//LLL7VgwQIVFRWpsbFRzz33nJ555hnt3bvX6VGBEGvXriW7MBLZRSyLs64+1KIDZs+erZkzZ2rLli39+/Ly8rRo0SKVlZWFrP/tb3+r/fv3q6mpqX9fSUmJ/vGPf6iurm7I++vq6pLX61VGRoZrzpbZ1tYW7RFCuO2s0uPHO3qw5//Zk08+qW3btvVfdjK70n/zGwgEXHNG3OGc4XW0HT16NNoj2PQd1dYNLly4oPz8/Khld9u2ba45q/Sjjz4a7RFCuO3nafr06dEeoV9vb6+OHz8+rN9/jv517+7uVkNDQ/+5jfr4/X7V1taGvU1dXV3I+gceeECfffaZvvvuu5D1wWBQXV1dtg2IlPvuu892OZLZlcgvIq+7u1sS2UVsc7S8nDlzRj09PUpPT7ftT09PV0dHR9jbdHR0hF1/+fJlnTlzJmR9WVmZvF5v/5aVlRW5B4AxLy0tzXY5ktmVyC8ir7OzUxLZRWwblfdVrn2ZzLKsQV86C7c+3H7pyvu6gUCgf2ttbY3AxMAVTmZXIr9wDtlFLHP0gwapqamKj48PafunTp0Kafl9MjIywq4fP368pkyZErLe4/HI4/FEbmjgKidPnrRdjmR2JfKLyOv7/A3ZRSxz9JWXhIQE+Xw+VVVV2fZXVVVp7ty5YW9TWFgYsv7QoUOaNWuWJkyY4NisQDgffPCB7TLZhdslJCRIIruIbY6/bVRaWqrXXntN27dvV1NTk1avXq2WlhaVlJRIuvLS4+OPP96/vqSkRF9//bVKS0vV1NSk7du3a9u2bVqzZo3TowIhdu7cSXZhJLKLWOb4/09dsmSJzp49qw0bNqi9vV0FBQU6cOCAsrOzJUnt7e22Yw/k5OTowIEDWr16tTZv3qypU6fq5Zdf1iOPPOL0qECIsrIysgsjkV3EMseP8zLaOM7L8HCcl+EZ7eOtcJyX4eE4LwPrO85LtLLLcV4G57afJ47zAgAAMAooLwAAwCiUFwAAYBTKCwAAMArlBQAAGIXyAgAAjEJ5AQAARhmV8lJeXq6cnBwlJibK5/PpyJEjA66trq5WXFxcyPb555+PxqiATUVFBdmFkcguYpnj5aWyslKrVq3SunXr1NjYqKKiIs2fP992dMdwmpub1d7e3r/l5uY6PSoQYu3atWQXRiK7iGWOl5eXXnpJS5cu1bJly5SXl6dNmzYpKytLW7ZsGfR2aWlpysjI6N/i4+OdHhUIUVxcTHZhJLKLWOboMdm7u7vV0NCgZ5991rbf7/ertrZ20Nvedttt+vbbb/WjH/1Iv/vd73TvvfeGXRcMBhUMBvsvBwIBSVcOM+wWXV1d0R4hhNtOD+BW1+YuktmVBs6vGzPjJhcvXoz2CDZu+iN//vx5SdHL7jfffPO/jh5x/BwNzU1/K/tmGdZZiywHtbW1WZKsv//977b9GzdutG666aawt/n888+trVu3Wg0NDVZtba3161//2oqLi7NqamrCrl+/fr0liY3NkW3fvn2OZZf8sjm5kV02U7fW1tYBc9fH0RMznjhxQjfccINqa2tVWFjYv3/jxo164403hv1hsIULFyouLk779+8Pue7a9t/b26tz585pypQp3/sEWF1dXcrKylJra6trTpLnNrH6HJ04cUJ5eXn66KOPdOedd/bvj2R2JefyG6vfl0iK1eeI7Ma+WH2OLMvShQsXNHXq1CFPrOzo20apqamKj49XR0eHbf+pU6eUnp4+7K8zZ84c7dq1K+x1Ho9HHo/Hti/SZ3hNTk6OqYA4Idaeo8TERMXHx+vUqVO2/ZHMruR8fmPt++KEWHuOyO7YEYvPkdfrHdY6Rz+wm5CQIJ/Pp6qqKtv+qqoqzZ07d9hfp7GxUZmZmZEeDxgQ2YWpyC7GAkdfeZGk0tJSFRcXa9asWSosLNTWrVvV0tKikpISSVf+O19bW5t27twpSdq0aZNmzJih/Px8dXd3a9euXdq7d6/27t3r9KiADdmFqcguYp3j5WXJkiU6e/asNmzYoPb2dhUUFOjAgQPKzs6WJLW3t9uOPdDd3a01a9aora1NSUlJys/P13vvvacFCxY4PWoIj8ej9evXh7w0iv+K5eeI7Ma2WH6OyG5s4zmSHP3ALgAAQKRxbiMAAGAUygsAADAK5QUAABiF8gIAAIxCeRlAeXn5sE8nP9aUlZXp9ttv16RJk5SWlqZFixapubk52mPhP8juwMiuu5HdgZFdO8pLGJWVlVq1atWITyc/VtTU1GjFihWqr69XVVWVLl++LL/fr0uXLkV7tDGP7A6O7LoX2R0c2bXjv0qHMXv2bM2cOdN2+vi8vDwtWrRIZWVlUZzMnU6fPq20tDTV1NTo7rvvjvY4YxrZHRmy6x5kd2TGenZ55eUa3d3damhokN/vt+0fzunkx6pAICBJmjx5cpQnGdvI7siRXXcguyM31rNLebnGmTNn1NPTE3ICs/T09JATTOLKWUBLS0t11113qaCgINrjjGlkd2TIrnuQ3ZEhu6NwegBTXXtKd8uyvtdp3mPVypUrdfToUX300UfRHgX/QXaHh+y6D9kdHrJLeQmRmpqq+Pj4kLY/0tPJjwVPP/209u/frw8//FDTpk2L9jhjHtkdPrLrLmR3+MjuFbxtdI1InU4+llmWpZUrV+qtt97S4cOHlZOTE+2RILI7HGTXncju0MiuHa+8hDHU6eTHuhUrVujNN9/U22+/rUmTJvX/a8nr9SopKSnK041tZHdwZNe9yO7gyO41LIS1efNmKzs720pISLBmzpxp1dTURHsk15AUdtuxY0e0R4NFdgdDdt2N7A6M7NpxnBcAAGAUPvMCAACMQnkBAABGobwAAACjUF4AAIBRKC8AAMAolBcAAGAUygsAADAK5QUAABiF8gIAAIxCeQEAAEahvAAAAKNQXgAAgFH+H8wC4ke/apfXAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 640x480 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "examples = iter(test_loader)\n",
    "example_data, example_targets = examples.next()\n",
    "for i in range(6):\n",
    "    plt.subplot(2, 3, i + 1)\n",
    "    if DATASET_NAME == \"MNIST\":\n",
    "        plt.imshow(example_data[i][0], cmap=\"gray\")\n",
    "    elif DATASET_NAME == \"FFOOM\":\n",
    "        plt.imshow(example_data[i].unsqueeze(dim=0), cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b73031b1-c43a-4d41-b6b8-70e80f7041ad",
   "metadata": {},
   "source": [
    "### Define training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a2b9e166",
   "metadata": {},
   "outputs": [],
   "source": [
    "def validation(\n",
    "    model,\n",
    "    optimizer,\n",
    "    val_loader,\n",
    "    l1_weight,\n",
    "    l2_weight,\n",
    "    epoch,\n",
    "    batch,\n",
    "    num_strata=10,\n",
    "    num_elements_per_strata=100,\n",
    "    visualize=False,\n",
    "):\n",
    "    model.eval()\n",
    "    # Construct stratified validation set based on the 0th feature value\n",
    "    val_inputs, val_labels = next(iter(val_loader))\n",
    "\n",
    "    sorted_row_indices = val_inputs[:, 0].argsort()\n",
    "    val_inputs_sorted = val_inputs[sorted_row_indices].numpy()\n",
    "    val_labels_sorted = val_labels[sorted_row_indices].numpy()\n",
    "\n",
    "    val_inputs_stratified = np.split(val_inputs_sorted, num_strata, axis=0)\n",
    "    val_labels_stratified = np.split(val_labels_sorted, num_strata, axis=0)\n",
    "\n",
    "    if num_elements_per_strata < len(val_inputs_stratified[0]):\n",
    "        sample_inds = np.random.choice(\n",
    "            len(val_inputs_stratified[0]), num_elements_per_strata, replace=False\n",
    "        )\n",
    "        print(sample_inds)\n",
    "        wandb.log({\"sample_inds_for_validation\": sample_inds})\n",
    "        val_inputs_stratified = [inp[sample_inds] for inp in val_inputs_stratified]\n",
    "        val_labels_stratified = [label[sample_inds] for label in val_labels_stratified]\n",
    "\n",
    "    total_correct, total_entropy, total_val_points = 0, 0, 0\n",
    "    total_entropy_per_feature = np.zeros(shape=(4,))\n",
    "\n",
    "    val_metrics_list = []\n",
    "    val_metrics = dict()\n",
    "\n",
    "    import tempfile\n",
    "\n",
    "    temp_dir = (\n",
    "        tempfile.TemporaryDirectory()\n",
    "    )  # for saving and logging graph visualizations\n",
    "\n",
    "    for i in range(num_strata):\n",
    "        inputs = Node(val_inputs_stratified[i], name=\"val_inputs\")\n",
    "        labels = val_labels_stratified[i]\n",
    "\n",
    "        # Apply model and compute entropy\n",
    "        outputs = model(inputs)\n",
    "        unregularized_loss = cross_entropy_loss(outputs, labels)\n",
    "        loss = unregularized_loss + regularize(\n",
    "            model=model, l1_weight=l1_weight, l2_weight=l2_weight\n",
    "        )\n",
    "\n",
    "        # Zero out all semiring values (e.g. gradients, entropy, etc) from the previous minibatch to compute the gradients and entropy correctly.\n",
    "        optimizer.zero_gradients()\n",
    "        loss.backprop(values_to_compute=(\"abs_val_grad\", \"entropy\", \"grad\"))\n",
    "        grads_per_input_feature = bf.reduce_mean(inputs.grad, axis=0)\n",
    "        abs_grads_per_input_feature = bf.reduce_mean(bf.abs(inputs.grad), axis=0)\n",
    "        entropy_per_input_feature = bf.reduce_mean(inputs.compute_entropy(), axis=0)\n",
    "        entropy = bf.reduce_mean(inputs.compute_entropy())\n",
    "\n",
    "        num_correct_in_strata = sum(np.argmax(outputs.val, axis=1) == labels)\n",
    "        accuracy_in_strata = num_correct_in_strata / len(labels)\n",
    "\n",
    "        total_correct += num_correct_in_strata\n",
    "        total_val_points += len(labels)\n",
    "        total_entropy += bf.reduce_sum(inputs.compute_entropy()).val\n",
    "        total_entropy_per_feature += bf.reduce_sum(inputs.compute_entropy(), axis=0).val\n",
    "\n",
    "        val_metrics[f\"strata_{i}\"] = {\n",
    "            \"unregularized_loss\": unregularized_loss.val,\n",
    "            \"loss\": loss.val,\n",
    "            \"entropy\": entropy.val,\n",
    "            \"val_entropy_per_input_feature\": entropy_per_input_feature.val,\n",
    "            \"val_grads_per_input_feature\": grads_per_input_feature.val,\n",
    "            \"val_abs_grads_per_input_feature\": abs_grads_per_input_feature.val,\n",
    "            \"epoch\": epoch,\n",
    "            \"batch\": batch,\n",
    "            \"accuracy\": accuracy_in_strata,\n",
    "        }\n",
    "\n",
    "        if visualize:\n",
    "            try:\n",
    "                loss.visualize(\n",
    "                    save_path=os.path.join(\n",
    "                        temp_dir.name, f\"strata_{i}_e{epoch}_b{batch}_visualization.png\"\n",
    "                    ),\n",
    "                    vals_to_include={\"entropy\"},\n",
    "                )\n",
    "            except Exception as e:\n",
    "                print(f\"Unable to visaulize loss node computation graph due to error: {e}\")\n",
    "\n",
    "    val_metrics[\"all\"] = {\n",
    "        \"accuracy\": total_correct / total_val_points,\n",
    "        \"epoch\": epoch,\n",
    "        \"batch\": batch,\n",
    "        \"entropy\": total_entropy / total_val_points / 4,  # average across all features\n",
    "        \"total_entropy_relevant\": total_entropy_per_feature[0] / total_val_points,\n",
    "        \"total_entropy_irrelevant\": np.mean(total_entropy_per_feature[1:])\n",
    "        / total_val_points,\n",
    "    }\n",
    "    \n",
    "    if wandb.run is not None:\n",
    "        wandb.log(val_metrics)\n",
    "\n",
    "    if visualize:\n",
    "        print(f\"Logging visualizations info to w&b run {wandb.run}.\")\n",
    "        artifact = wandb.Artifact(\n",
    "            name=f\"{data_id}_visualizations\", type=\"visualization\"\n",
    "        )\n",
    "        artifact.add_dir(local_path=temp_dir.name)\n",
    "        run.log_artifact(artifact)\n",
    "\n",
    "    temp_dir.cleanup()\n",
    "    model.train()\n",
    "\n",
    "    return val_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "731a119a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(\n",
    "    model,\n",
    "    optimizer,\n",
    "    train_loader,\n",
    "    model_dir,\n",
    "    overwrite_model=False,\n",
    "    num_epochs=NUM_EPOCHS,\n",
    "    l1_weight=0,\n",
    "    l2_weight=0,\n",
    "):\n",
    "    model.train()\n",
    "    os.makedirs(model_dir, exist_ok=True)\n",
    "    model_path = os.path.join(model_dir, \"model.bf\")\n",
    "\n",
    "    # Load model if already trained\n",
    "    if os.path.isfile(model_path):\n",
    "        if overwrite_model:\n",
    "            print(f\"Retraining and overwriting model at {model_path}.\")\n",
    "        else:\n",
    "            print(f\"Loading trained model from {model_path}.\")\n",
    "            with open(model_path, \"rb\") as f:\n",
    "                model = pickle.load(f)\n",
    "            # Wow there's some DANGER here of the optimizer being attached to a different model than the one being loaded here...oof\n",
    "\n",
    "            return model\n",
    "\n",
    "    # Train the model\n",
    "    entropies = dict()\n",
    "    losses = dict()\n",
    "    n_total_steps = len(train_loader)\n",
    "    for epoch in range(num_epochs):\n",
    "        for i, (inputs, labels) in enumerate(train_loader):\n",
    "            inputs = Node(inputs.numpy(), name=\"input_batch\")\n",
    "            labels = labels.numpy()\n",
    "\n",
    "            # Forward pass\n",
    "            outputs = model(inputs)\n",
    "            unregularized_loss = cross_entropy_loss(outputs, labels)\n",
    "            loss = unregularized_loss + regularize(\n",
    "                model=model, l1_weight=l1_weight, l2_weight=l2_weight\n",
    "            )\n",
    "\n",
    "            # Zero out all semiring values (e.g. gradients, entropy, etc) from the previous minibatch to compute the gradients and entropy correctly.\n",
    "            optimizer.zero_gradients()\n",
    "            unregularized_loss.backprop(\n",
    "                values_to_compute=(\"abs_val_grad\", \"entropy\", \"grad\")\n",
    "            )\n",
    "\n",
    "            grads_per_input_feature = bf.reduce_mean(inputs.grad, axis=0)\n",
    "            entropy_per_input_feature = bf.reduce_mean(inputs.compute_entropy(), axis=0)\n",
    "            entropy = bf.reduce_mean(inputs.compute_entropy())\n",
    "\n",
    "            optimizer.step()\n",
    "\n",
    "            if wandb.run is not None:\n",
    "                wandb.log(\n",
    "                    {\n",
    "                        \"unregularized_loss\": unregularized_loss.val,\n",
    "                        \"loss\": loss.val,\n",
    "                        \"entropy\": entropy.val,\n",
    "                        \"entropy_per_input_feature\": entropy_per_input_feature.val,\n",
    "                        \"grads_per_input_feature\": grads_per_input_feature.val,\n",
    "                        \"epoch\": epoch,\n",
    "                        \"batch\": i,\n",
    "                    }\n",
    "                )\n",
    "\n",
    "            if i % 100 == 0:\n",
    "                print(\n",
    "                    f\"Epoch [{epoch+1}/{NUM_EPOCHS}], Step [{i}/{n_total_steps}], \"\n",
    "                    # f'Total Loss: {regularized_loss.val:.4f}, '\n",
    "                    f\"Unregularized Cross-Entropy Loss: {unregularized_loss.val:.4f}, \"\n",
    "                    f\"Cross-Entropy Loss: {loss.val:.4f}, \"\n",
    "                    f\"Entropy: {entropy.val:.4f}\"\n",
    "                )\n",
    "\n",
    "                validation(\n",
    "                    model,\n",
    "                    optimizer,\n",
    "                    val_loader,\n",
    "                    l1_weight=l1_weight,\n",
    "                    l2_weight=l2_weight,\n",
    "                    epoch=epoch,\n",
    "                    batch=i,\n",
    "                    num_strata=10,\n",
    "                    visualize=False,\n",
    "                ),\n",
    "\n",
    "                train_step_num = i + n_total_steps * epoch\n",
    "                entropies[train_step_num] = entropy.val\n",
    "                losses[train_step_num] = loss.val\n",
    "\n",
    "                # Construct and save entropies and losses to CSV\n",
    "                entropies_df = pd.DataFrame([entropies]).melt(\n",
    "                    var_name=\"steps\", value_name=\"entropy\"\n",
    "                )\n",
    "                losses_df = pd.DataFrame([losses]).melt(\n",
    "                    var_name=\"steps\", value_name=\"loss\"\n",
    "                )\n",
    "\n",
    "                entropies_df.to_csv(os.path.join(model_dir, \"entropies.csv\"))\n",
    "                losses_df.to_csv(os.path.join(model_dir, \"losses.csv\"))\n",
    "\n",
    "    # Save model\n",
    "    with open(model_path, \"wb\") as f:\n",
    "        pickle.dump(model, f)\n",
    "\n",
    "    # Log to W&B\n",
    "    if wandb.run is not None:\n",
    "        print(f\"Logging model to w&b run {wandb.run}.\")\n",
    "        artifact = wandb.Artifact(name=\"model\", type=\"model\")\n",
    "        artifact.add_file(local_path=model_path)\n",
    "        run.log_artifact(artifact)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcf1982a-c924-45bb-87ec-14d67db0d8c6",
   "metadata": {},
   "source": [
    "### Define evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1cf416a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(model, model_dir, test_loader, overwrite_metrics=False):\n",
    "    model.eval()\n",
    "    metrics_path = os.path.join(model_dir, \"metrics.csv\")\n",
    "    if os.path.isfile(metrics_path):\n",
    "        if overwrite_metrics:\n",
    "            print(f\"Recomputing and overwriting metrics at {metrics_path}.\")\n",
    "        else:\n",
    "            print(f\"Loading precomputed metrics from {metrics_path}.\")\n",
    "            return pd.read_csv(metrics_path)\n",
    "\n",
    "    metrics = dict()\n",
    "    total_correct = 0\n",
    "    total_test_points = 0\n",
    "    for inputs, labels in test_loader:\n",
    "        inputs = inputs.numpy()\n",
    "        labels = labels.numpy()\n",
    "        outputs = model(inputs)\n",
    "        num_correct_in_batch = sum(np.argmax(outputs.val, axis=1) == labels)\n",
    "        total_correct += num_correct_in_batch\n",
    "        total_test_points += len(labels)\n",
    "\n",
    "    accuracy = total_correct / total_test_points\n",
    "\n",
    "    metrics[\"accuracy\"] = accuracy\n",
    "    metrics_df = pd.DataFrame.from_dict([metrics])\n",
    "\n",
    "    metrics_df.to_csv(os.path.join(model_dir, \"metrics.csv\"))\n",
    "\n",
    "    # Log to W&B\n",
    "    if wandb.run is not None:\n",
    "        print(f\"Logging metrics to w&b run {wandb.run}.\")\n",
    "        wandb.log(metrics)\n",
    "\n",
    "    return metrics_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be668f7e-005f-4850-9052-7e544a0e51a0",
   "metadata": {},
   "source": [
    "### Train and evaluate models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8a105830",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile_dir = os.path.join(model_dir, \"mem_profiles\")\n",
    "os.makedirs(profile_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "026b6bc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time: 2.102 seconds\n",
      "Retraining and overwriting model at data/FFOOM/FFOOM-num_points10000/1.0-2221/models/MLP-hidden_sz8-bs32-lr0.001-n1/model.bf.\n",
      "Epoch [1/1], Step [0/219], Unregularized Cross-Entropy Loss: 0.8982, Cross-Entropy Loss: 0.8982, Entropy: 2.3416\n",
      "2.3656337 2.365634\n",
      "Epoch [1/1], Step [100/219], Unregularized Cross-Entropy Loss: 0.7549, Cross-Entropy Loss: 0.7549, Entropy: 2.5746\n",
      "2.5698583 2.5698583\n",
      "Epoch [1/1], Step [200/219], Unregularized Cross-Entropy Loss: 0.6146, Cross-Entropy Loss: 0.6146, Entropy: 2.7648\n",
      "2.723858 2.723858\n",
      "Logging model to w&b run <wandb.sdk.wandb_run.Run object at 0x7fa4132f5c70>.\n",
      "Time: 44.408 seconds\n",
      "Recomputing and overwriting metrics at data/FFOOM/FFOOM-num_points10000/1.0-2221/models/MLP-hidden_sz8-bs32-lr0.001-n1/metrics.csv.\n",
      "Logging metrics to w&b run <wandb.sdk.wandb_run.Run object at 0x7fa4132f5c70>.\n",
      "Time: 0.224 seconds\n",
      "     accuracy\n",
      "0  0.59000003\n"
     ]
    }
   ],
   "source": [
    "# Initialize model\n",
    "with catchtime() as t:\n",
    "    model = MLP(\n",
    "        input_size=INPUT_SIZE,\n",
    "        hidden_size=HIDDEN_SIZE,\n",
    "        num_classes=NUM_CLASSES,\n",
    "        dropout_prob=DROPOUT_PROB,\n",
    "    )\n",
    "init_model_time = t.time\n",
    "init_model_nvidia_mem_used = gpu_memory_usage()\n",
    "try:\n",
    "    jax.profiler.save_device_memory_profile(\n",
    "        os.path.join(\n",
    "            profile_dir, f\"randn{INPUT_SIZE}x{HIDDEN_SIZE}x{NUM_CLASSES}_mlp.prof\"\n",
    "        )\n",
    "    )\n",
    "except:\n",
    "    print(\"Can't jax profile because no jax\")\n",
    "\n",
    "# Initialize optimizer\n",
    "optimizer = Adam(model.parameters(), step_size=LEARNING_RATE)\n",
    "\n",
    "# Train model\n",
    "with catchtime() as t:\n",
    "    model = train_model(\n",
    "        model,\n",
    "        optimizer,\n",
    "        train_loader,\n",
    "        model_dir,\n",
    "        overwrite_model=OVERWRITE_MODEL,\n",
    "        num_epochs=NUM_EPOCHS,\n",
    "        l1_weight=L1_WEIGHT,\n",
    "        l2_weight=L2_WEIGHT,\n",
    "    )\n",
    "train_model_time = t.time\n",
    "train_model_nvidia_mem_used = gpu_memory_usage()\n",
    "try:\n",
    "    jax.profiler.save_device_memory_profile(\n",
    "        os.path.join(\n",
    "            profile_dir,\n",
    "            f\"randn{INPUT_SIZE}x{HIDDEN_SIZE}x{NUM_CLASSES}_train_model.prof\",\n",
    "        )\n",
    "    )\n",
    "except:\n",
    "    print(\"Can't jax profile because no jax\")\n",
    "\n",
    "# Evaluate model\n",
    "with catchtime() as t:\n",
    "    metrics = evaluate_model(\n",
    "        model, model_dir, test_loader, overwrite_metrics=OVERWRITE_MODEL\n",
    "    )\n",
    "eval_model_time = t.time\n",
    "eval_model_nvidia_mem_used = gpu_memory_usage()\n",
    "try:\n",
    "    jax.profiler.save_device_memory_profile(\n",
    "        os.path.join(\n",
    "            profile_dir,\n",
    "            f\"randn{INPUT_SIZE}x{HIDDEN_SIZE}x{NUM_CLASSES}_eval_model.prof\",\n",
    "        )\n",
    "    )\n",
    "except:\n",
    "    print(\"Can't jax profile because no jax\")\n",
    "print(metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "aa5d98a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.log(\n",
    "    {\n",
    "        \"init_model_time\": init_model_time,\n",
    "        \"train_model_time\": train_model_time,\n",
    "        \"eval_model_time\": eval_model_time,\n",
    "        \"init_model_nvidia_mem_used\": init_model_nvidia_mem_used,\n",
    "        \"eval_model_nvidia_mem_used\": eval_model_nvidia_mem_used,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "79746a4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Adding directory to artifact (./data/FFOOM/FFOOM-num_points10000/1.0-2221/models/MLP-hidden_sz8-bs32-lr0.001-n1/mem_profiles)... Done. 0.0s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logging profiling info to w&b run <wandb.sdk.wandb_run.Run object at 0x7fa4132f5c70>.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<wandb.sdk.wandb_artifacts.Artifact at 0x7fa32c18fd90>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"Logging profiling info to w&b run {wandb.run}.\")\n",
    "artifact = wandb.Artifact(name=f\"{data_id}_profiles\", type=\"profiles\")\n",
    "artifact.add_dir(local_path=profile_dir)\n",
    "run.log_artifact(artifact)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5c7a3d57",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[68 46 97  7 65 55 74 83 11 21 29 31 88 75 47 56 80  6 69 20 44 39 10 78\n",
      " 66 59 51 41 62 64 34 70]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Adding directory to artifact (/tmp/tmpm3rsp6h0)... Done. 0.0s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.703433 2.7034333\n",
      "Logging visualizations info to w&b run <wandb.sdk.wandb_run.Run object at 0x7fa4132f5c70>.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'strata_0': {'unregularized_loss': DeviceArray(0.7161944, dtype=float32),\n",
       "  'loss': DeviceArray(0.7161944, dtype=float32),\n",
       "  'entropy': DeviceArray(2.6115732, dtype=float32),\n",
       "  'val_entropy_per_input_feature': DeviceArray([2.7091613, 2.582757 , 2.6365573, 2.5178163], dtype=float32),\n",
       "  'val_grads_per_input_feature': DeviceArray([ 0.004101  , -0.00481713, -0.00242533, -0.00686217], dtype=float32),\n",
       "  'val_abs_grads_per_input_feature': DeviceArray([0.004101  , 0.00481713, 0.00242533, 0.00686217], dtype=float32),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'accuracy': DeviceArray(0.375, dtype=float32, weak_type=True)},\n",
       " 'strata_1': {'unregularized_loss': DeviceArray(0.73381954, dtype=float32),\n",
       "  'loss': DeviceArray(0.73381954, dtype=float32),\n",
       "  'entropy': DeviceArray(2.6459196, dtype=float32),\n",
       "  'val_entropy_per_input_feature': DeviceArray([2.7131646, 2.642682 , 2.661244 , 2.5665872], dtype=float32),\n",
       "  'val_grads_per_input_feature': DeviceArray([ 0.00682917, -0.00407612, -0.00306033, -0.00652896], dtype=float32),\n",
       "  'val_abs_grads_per_input_feature': DeviceArray([0.00682917, 0.00407612, 0.00306033, 0.00652896], dtype=float32),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'accuracy': DeviceArray(0.34375, dtype=float32, weak_type=True)},\n",
       " 'strata_2': {'unregularized_loss': DeviceArray(0.7410941, dtype=float32),\n",
       "  'loss': DeviceArray(0.7410941, dtype=float32),\n",
       "  'entropy': DeviceArray(2.7236748, dtype=float32),\n",
       "  'val_entropy_per_input_feature': DeviceArray([2.6912394, 2.8002002, 2.710739 , 2.6925192], dtype=float32),\n",
       "  'val_grads_per_input_feature': DeviceArray([ 0.01423473, -0.00196485, -0.00466619, -0.00687591], dtype=float32),\n",
       "  'val_abs_grads_per_input_feature': DeviceArray([0.01423473, 0.00196485, 0.00466619, 0.00687591], dtype=float32),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'accuracy': DeviceArray(0.28125, dtype=float32, weak_type=True)},\n",
       " 'strata_3': {'unregularized_loss': DeviceArray(0.83091867, dtype=float32),\n",
       "  'loss': DeviceArray(0.83091867, dtype=float32),\n",
       "  'entropy': DeviceArray(2.731389, dtype=float32),\n",
       "  'val_entropy_per_input_feature': DeviceArray([2.6825569, 2.8351276, 2.6946812, 2.713191 ], dtype=float32),\n",
       "  'val_grads_per_input_feature': DeviceArray([ 0.01649533, -0.0019092 , -0.00523379, -0.008914  ], dtype=float32),\n",
       "  'val_abs_grads_per_input_feature': DeviceArray([0.01649533, 0.0019092 , 0.00523379, 0.008914  ], dtype=float32),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'accuracy': DeviceArray(0.0625, dtype=float32, weak_type=True)},\n",
       " 'strata_4': {'unregularized_loss': DeviceArray(0.87453294, dtype=float32),\n",
       "  'loss': DeviceArray(0.87453294, dtype=float32),\n",
       "  'entropy': DeviceArray(2.7254274, dtype=float32),\n",
       "  'val_entropy_per_input_feature': DeviceArray([2.6625197, 2.8381577, 2.684137 , 2.7168949], dtype=float32),\n",
       "  'val_grads_per_input_feature': DeviceArray([ 0.01745269, -0.00178183, -0.00551826, -0.00850037], dtype=float32),\n",
       "  'val_abs_grads_per_input_feature': DeviceArray([0.01745269, 0.00178183, 0.00551826, 0.00850037], dtype=float32),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'accuracy': DeviceArray(0., dtype=float32, weak_type=True)},\n",
       " 'strata_5': {'unregularized_loss': DeviceArray(0.5344088, dtype=float32),\n",
       "  'loss': DeviceArray(0.5344088, dtype=float32),\n",
       "  'entropy': DeviceArray(2.7290506, dtype=float32),\n",
       "  'val_entropy_per_input_feature': DeviceArray([2.6148245, 2.8746207, 2.6476066, 2.7791505], dtype=float32),\n",
       "  'val_grads_per_input_feature': DeviceArray([-0.01002192,  0.00120268,  0.0031193 ,  0.00471677], dtype=float32),\n",
       "  'val_abs_grads_per_input_feature': DeviceArray([0.01232255, 0.00142099, 0.00385546, 0.005702  ], dtype=float32),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'accuracy': DeviceArray(0.9375, dtype=float32, weak_type=True)},\n",
       " 'strata_6': {'unregularized_loss': DeviceArray(0.48096982, dtype=float32),\n",
       "  'loss': DeviceArray(0.48096982, dtype=float32),\n",
       "  'entropy': DeviceArray(2.7239866, dtype=float32),\n",
       "  'val_entropy_per_input_feature': DeviceArray([2.6071398, 2.8723826, 2.6392283, 2.7771966], dtype=float32),\n",
       "  'val_grads_per_input_feature': DeviceArray([-0.01145303,  0.00133694,  0.0035782 ,  0.00531129], dtype=float32),\n",
       "  'val_abs_grads_per_input_feature': DeviceArray([0.01145303, 0.00133694, 0.0035782 , 0.00531129], dtype=float32),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'accuracy': DeviceArray(1., dtype=float32, weak_type=True)},\n",
       " 'strata_7': {'unregularized_loss': DeviceArray(0.43408626, dtype=float32),\n",
       "  'loss': DeviceArray(0.43408626, dtype=float32),\n",
       "  'entropy': DeviceArray(2.7117066, dtype=float32),\n",
       "  'val_entropy_per_input_feature': DeviceArray([2.596474 , 2.8527436, 2.6131575, 2.7844508], dtype=float32),\n",
       "  'val_grads_per_input_feature': DeviceArray([-0.01061929,  0.00117135,  0.00331605,  0.00572638], dtype=float32),\n",
       "  'val_abs_grads_per_input_feature': DeviceArray([0.01061929, 0.00117135, 0.00331605, 0.00572638], dtype=float32),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'accuracy': DeviceArray(1., dtype=float32, weak_type=True)},\n",
       " 'strata_8': {'unregularized_loss': DeviceArray(0.3899991, dtype=float32),\n",
       "  'loss': DeviceArray(0.3899991, dtype=float32),\n",
       "  'entropy': DeviceArray(2.7191622, dtype=float32),\n",
       "  'val_entropy_per_input_feature': DeviceArray([2.6048462, 2.862568 , 2.6209302, 2.7883034], dtype=float32),\n",
       "  'val_grads_per_input_feature': DeviceArray([-0.00989616,  0.00126748,  0.0031149 ,  0.00497575], dtype=float32),\n",
       "  'val_abs_grads_per_input_feature': DeviceArray([0.00989616, 0.00126748, 0.0031149 , 0.00497575], dtype=float32),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'accuracy': DeviceArray(1., dtype=float32, weak_type=True)},\n",
       " 'strata_9': {'unregularized_loss': DeviceArray(0.36122996, dtype=float32),\n",
       "  'loss': DeviceArray(0.36122996, dtype=float32),\n",
       "  'entropy': DeviceArray(2.712441, dtype=float32),\n",
       "  'val_entropy_per_input_feature': DeviceArray([2.5958822, 2.8605933, 2.6195078, 2.7737803], dtype=float32),\n",
       "  'val_grads_per_input_feature': DeviceArray([-0.00925637,  0.0012126 ,  0.0029164 ,  0.00429633], dtype=float32),\n",
       "  'val_abs_grads_per_input_feature': DeviceArray([0.00925637, 0.0012126 , 0.0029164 , 0.00429633], dtype=float32),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'accuracy': DeviceArray(1., dtype=float32, weak_type=True)},\n",
       " 'all': {'accuracy': DeviceArray(0.6, dtype=float32, weak_type=True),\n",
       "  'epoch': 1,\n",
       "  'batch': 5,\n",
       "  'entropy': DeviceArray(2.703433, dtype=float32),\n",
       "  'total_entropy_relevant': DeviceArray(2.6477811, dtype=float32),\n",
       "  'total_entropy_irrelevant': DeviceArray(2.721984, dtype=float32)}}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation(\n",
    "    model,\n",
    "    optimizer,\n",
    "    val_loader,\n",
    "    l1_weight=L1_WEIGHT,\n",
    "    l2_weight=L2_WEIGHT,\n",
    "    epoch=1,\n",
    "    batch=i,\n",
    "    num_strata=10,\n",
    "    num_elements_per_strata=NUM_ELEMENTS_PER_STRATA,\n",
    "    visualize=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "47ebc592",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>accuracy</td><td>▁</td></tr><tr><td>batch</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███</td></tr><tr><td>entropy</td><td>▂▂▂▁▁▁▃▁▂▂▂▂▃▄▅▅▅▄▆▆▆▇▇██▇███████████▇█▇</td></tr><tr><td>epoch</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>eval_model_time</td><td>▁</td></tr><tr><td>init_model_time</td><td>▁</td></tr><tr><td>loss</td><td>▃▄█▆▆▃▅▃▅▆▃▆▃▄▄▃▄▄▄▃▃▃▃▂▂▂▂▃▂▂▃▂▂▁▂▁▂▂▁▂</td></tr><tr><td>train_model_time</td><td>▁</td></tr><tr><td>unregularized_loss</td><td>▃▄█▆▆▃▅▃▅▆▃▆▃▄▄▃▄▄▄▃▃▃▃▂▂▂▂▃▂▂▃▂▂▁▂▁▂▂▁▂</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>accuracy</td><td>0.59</td></tr><tr><td>batch</td><td>218</td></tr><tr><td>entropy</td><td>2.71827</td></tr><tr><td>epoch</td><td>0</td></tr><tr><td>eval_model_nvidia_mem_used</td><td>1963 MiB</td></tr><tr><td>eval_model_time</td><td>0.22354</td></tr><tr><td>init_model_nvidia_mem_used</td><td>1465 MiB</td></tr><tr><td>init_model_time</td><td>2.10226</td></tr><tr><td>loss</td><td>0.64589</td></tr><tr><td>train_model_time</td><td>44.40837</td></tr><tr><td>unregularized_loss</td><td>0.64589</td></tr></table><br/></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Synced <strong style=\"color:#cdcd00\">glistening-monkey-628</strong>: <a href=\"https://wandb.ai/kdu/bauer-ffoom/runs/2u32blit\" target=\"_blank\">https://wandb.ai/kdu/bauer-ffoom/runs/2u32blit</a><br/>Synced 7 W&B file(s), 0 media file(s), 9 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20230203_192813-2u32blit/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d12306",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.15 ('jax')",
   "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.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "e140c34e0a267ffc2d97632a88da9d0788950c86e37f3cbf599f33bfe9d6bd4a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}