example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TabM\n",
    "\n",
    "This is a standalone usage example for the TabM project.\n",
    "The easiest way to run it is [Pixi](https://pixi.sh/latest/#installation):\n",
    "\n",
    "```shell\n",
    "git clone https://github.com/yandex-research/tabm\n",
    "cd tabm\n",
    "\n",
    "# With GPU:\n",
    "pixi run -e cuda jupyter-lab example.ipynb\n",
    "\n",
    "# Without GPU:\n",
    "pixi run jupyter-lab example.ipynb\n",
    "```\n",
    "\n",
    "For the full overview of the project, and for non-Pixi environment setups, see README in the repository:\n",
    "https://github.com/yandex-research/tabm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ruff: noqa: E402\n",
    "import math\n",
    "import random\n",
    "import warnings\n",
    "from typing import Literal, NamedTuple\n",
    "\n",
    "import numpy as np\n",
    "import rtdl_num_embeddings  # https://github.com/yandex-research/rtdl-num-embeddings\n",
    "import scipy.special\n",
    "import sklearn.datasets\n",
    "import sklearn.metrics\n",
    "import sklearn.model_selection\n",
    "import sklearn.preprocessing\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.optim\n",
    "from torch import Tensor\n",
    "from tqdm.std import tqdm\n",
    "\n",
    "warnings.simplefilter('ignore')\n",
    "from tabm_reference import Model, make_parameter_groups\n",
    "\n",
    "warnings.resetwarnings()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 0\n",
    "random.seed(seed)\n",
    "np.random.seed(seed + 1)\n",
    "torch.manual_seed(seed + 2)\n",
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# >>> Dataset.\n",
    "TaskType = Literal['regression', 'binclass', 'multiclass']\n",
    "\n",
    "# Regression.\n",
    "task_type: TaskType = 'regression'\n",
    "n_classes = None\n",
    "dataset = sklearn.datasets.fetch_california_housing()\n",
    "X_cont: np.ndarray = dataset['data']\n",
    "Y: np.ndarray = dataset['target']\n",
    "\n",
    "# Classification.\n",
    "# n_classes = 2\n",
    "# assert n_classes >= 2\n",
    "# task_type: TaskType = 'binclass' if n_classes == 2 else 'multiclass'\n",
    "# X_cont, Y = sklearn.datasets.make_classification(\n",
    "#     n_samples=20000,\n",
    "#     n_features=8,\n",
    "#     n_classes=n_classes,\n",
    "#     n_informative=3,\n",
    "#     n_redundant=2,\n",
    "# )\n",
    "\n",
    "task_is_regression = task_type == 'regression'\n",
    "\n",
    "# >>> Continuous features.\n",
    "X_cont: np.ndarray = X_cont.astype(np.float32)\n",
    "n_cont_features = X_cont.shape[1]\n",
    "\n",
    "# >>> Categorical features.\n",
    "# NOTE: the above datasets do not have categorical features, however,\n",
    "# for the demonstration purposes, it is possible to generate them.\n",
    "cat_cardinalities = [\n",
    "    # NOTE: uncomment the two lines below to add two categorical features.\n",
    "    # 4,  # Allowed values: [0, 1, 2, 3].\n",
    "    # 7,  # Allowed values: [0, 1, 2, 3, 4, 5, 6].\n",
    "]\n",
    "X_cat = (\n",
    "    np.column_stack(\n",
    "        [np.random.randint(0, c, (len(X_cont),)) for c in cat_cardinalities]\n",
    "    )\n",
    "    if cat_cardinalities\n",
    "    else None\n",
    ")\n",
    "\n",
    "# >>> Labels.\n",
    "if task_type == 'regression':\n",
    "    Y = Y.astype(np.float32)\n",
    "else:\n",
    "    assert n_classes is not None\n",
    "    Y = Y.astype(np.int64)\n",
    "    assert set(Y.tolist()) == set(\n",
    "        range(n_classes)\n",
    "    ), 'Classification labels must form the range [0, 1, ..., n_classes - 1]'\n",
    "\n",
    "# >>> Split the dataset.\n",
    "all_idx = np.arange(len(Y))\n",
    "trainval_idx, test_idx = sklearn.model_selection.train_test_split(\n",
    "    all_idx, train_size=0.8\n",
    ")\n",
    "train_idx, val_idx = sklearn.model_selection.train_test_split(\n",
    "    trainval_idx, train_size=0.8\n",
    ")\n",
    "data_numpy = {\n",
    "    'train': {'x_cont': X_cont[train_idx], 'y': Y[train_idx]},\n",
    "    'val': {'x_cont': X_cont[val_idx], 'y': Y[val_idx]},\n",
    "    'test': {'x_cont': X_cont[test_idx], 'y': Y[test_idx]},\n",
    "}\n",
    "if X_cat is not None:\n",
    "    data_numpy['train']['x_cat'] = X_cat[train_idx]\n",
    "    data_numpy['val']['x_cat'] = X_cat[val_idx]\n",
    "    data_numpy['test']['x_cat'] = X_cat[test_idx]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Feature preprocessing.\n",
    "# NOTE\n",
    "# The choice between preprocessing strategies depends on a task and a model.\n",
    "\n",
    "# Simple preprocessing strategy.\n",
    "# preprocessing = sklearn.preprocessing.StandardScaler().fit(\n",
    "#     data_numpy['train']['x_cont']\n",
    "# )\n",
    "\n",
    "# Advanced preprocessing strategy.\n",
    "# The noise is added to improve the output of QuantileTransformer in some cases.\n",
    "X_cont_train_numpy = data_numpy['train']['x_cont']\n",
    "noise = (\n",
    "    np.random.default_rng(0)\n",
    "    .normal(0.0, 1e-5, X_cont_train_numpy.shape)\n",
    "    .astype(X_cont_train_numpy.dtype)\n",
    ")\n",
    "preprocessing = sklearn.preprocessing.QuantileTransformer(\n",
    "    n_quantiles=max(min(len(train_idx) // 30, 1000), 10),\n",
    "    output_distribution='normal',\n",
    "    subsample=10**9,\n",
    ").fit(X_cont_train_numpy + noise)\n",
    "del X_cont_train_numpy\n",
    "\n",
    "# Apply the preprocessing.\n",
    "for part in data_numpy:\n",
    "    data_numpy[part]['x_cont'] = preprocessing.transform(data_numpy[part]['x_cont'])\n",
    "\n",
    "\n",
    "# Label preprocessing.\n",
    "class RegressionLabelStats(NamedTuple):\n",
    "    mean: float\n",
    "    std: float\n",
    "\n",
    "\n",
    "Y_train = data_numpy['train']['y'].copy()\n",
    "if task_type == 'regression':\n",
    "    # For regression tasks, it is highly recommended to standardize the training labels.\n",
    "    regression_label_stats = RegressionLabelStats(\n",
    "        Y_train.mean().item(), Y_train.std().item()\n",
    "    )\n",
    "    Y_train = (Y_train - regression_label_stats.mean) / regression_label_stats.std\n",
    "else:\n",
    "    regression_label_stats = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  PyTorch settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Device:        CUDA\n",
      "AMP:           False (dtype: torch.bfloat16)\n",
      "torch.compile: False\n"
     ]
    }
   ],
   "source": [
    "# Device\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Convert data to tensors\n",
    "data = {\n",
    "    part: {k: torch.as_tensor(v, device=device) for k, v in data_numpy[part].items()}\n",
    "    for part in data_numpy\n",
    "}\n",
    "Y_train = torch.as_tensor(Y_train, device=device)\n",
    "if task_type == 'regression':\n",
    "    for part in data:\n",
    "        data[part]['y'] = data[part]['y'].float()\n",
    "    Y_train = Y_train.float()\n",
    "\n",
    "# Automatic mixed precision (AMP)\n",
    "# torch.float16 is implemented for completeness,\n",
    "# but it was not tested in the project,\n",
    "# so torch.bfloat16 is used by default.\n",
    "amp_dtype = (\n",
    "    torch.bfloat16\n",
    "    if torch.cuda.is_available() and torch.cuda.is_bf16_supported()\n",
    "    else torch.float16\n",
    "    if torch.cuda.is_available()\n",
    "    else None\n",
    ")\n",
    "# Changing False to True will result in faster training on compatible hardware.\n",
    "amp_enabled = False and amp_dtype is not None\n",
    "grad_scaler = torch.cuda.amp.GradScaler() if amp_dtype is torch.float16 else None  # type: ignore\n",
    "\n",
    "# torch.compile\n",
    "compile_model = False\n",
    "\n",
    "# fmt: off\n",
    "print(\n",
    "    f'Device:        {device.type.upper()}'\n",
    "    f'\\nAMP:           {amp_enabled} (dtype: {amp_dtype})'\n",
    "    f'\\ntorch.compile: {compile_model}'\n",
    ")\n",
    "# fmt: on"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose one of the two configurations below.\n",
    "\n",
    "# TabM\n",
    "arch_type = 'tabm'\n",
    "bins = None\n",
    "\n",
    "# TabM-mini with the piecewise-linear embeddings.\n",
    "# arch_type = 'tabm-mini'\n",
    "# bins = rtdl_num_embeddings.compute_bins(data['train']['x_cont'])\n",
    "\n",
    "model = Model(\n",
    "    n_num_features=n_cont_features,\n",
    "    cat_cardinalities=cat_cardinalities,\n",
    "    n_classes=n_classes,\n",
    "    backbone={\n",
    "        'type': 'MLP',\n",
    "        'n_blocks': 3 if bins is None else 2,\n",
    "        'd_block': 512,\n",
    "        'dropout': 0.1,\n",
    "    },\n",
    "    bins=bins,\n",
    "    num_embeddings=(\n",
    "        None\n",
    "        if bins is None\n",
    "        else {\n",
    "            'type': 'PiecewiseLinearEmbeddings',\n",
    "            'd_embedding': 16,\n",
    "            'activation': False,\n",
    "            'version': 'B',\n",
    "        }\n",
    "    ),\n",
    "    arch_type=arch_type,\n",
    "    k=32,\n",
    ").to(device)\n",
    "optimizer = torch.optim.AdamW(make_parameter_groups(model), lr=2e-3, weight_decay=3e-4)\n",
    "\n",
    "if compile_model:\n",
    "    # NOTE\n",
    "    # `torch.compile` is intentionally called without the `mode` argument\n",
    "    # (mode=\"reduce-overhead\" caused issues during training with torch==2.0.1).\n",
    "    model = torch.compile(model)\n",
    "    evaluation_mode = torch.no_grad\n",
    "else:\n",
    "    evaluation_mode = torch.inference_mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test score before training: -1.1469\n"
     ]
    }
   ],
   "source": [
    "@torch.autocast(device.type, enabled=amp_enabled, dtype=amp_dtype)  # type: ignore[code]\n",
    "def apply_model(part: str, idx: Tensor) -> Tensor:\n",
    "    return (\n",
    "        model(\n",
    "            data[part]['x_cont'][idx],\n",
    "            data[part]['x_cat'][idx] if 'x_cat' in data[part] else None,\n",
    "        )\n",
    "        .squeeze(-1)  # Remove the last dimension for regression tasks.\n",
    "        .float()\n",
    "    )\n",
    "\n",
    "\n",
    "base_loss_fn = F.mse_loss if task_type == 'regression' else F.cross_entropy\n",
    "\n",
    "\n",
    "def loss_fn(y_pred: Tensor, y_true: Tensor) -> Tensor:\n",
    "    # TabM produces k predictions per object. Each of them must be trained separately.\n",
    "    # (regression)     y_pred.shape == (batch_size, k)\n",
    "    # (classification) y_pred.shape == (batch_size, k, n_classes)\n",
    "    k = y_pred.shape[-1 if task_type == 'regression' else -2]\n",
    "    return base_loss_fn(y_pred.flatten(0, 1), y_true.repeat_interleave(k))\n",
    "\n",
    "\n",
    "@evaluation_mode()\n",
    "def evaluate(part: str) -> float:\n",
    "    model.eval()\n",
    "\n",
    "    # When using torch.compile, you may need to reduce the evaluation batch size.\n",
    "    eval_batch_size = 8096\n",
    "    y_pred: np.ndarray = (\n",
    "        torch.cat(\n",
    "            [\n",
    "                apply_model(part, idx)\n",
    "                for idx in torch.arange(len(data[part]['y']), device=device).split(\n",
    "                    eval_batch_size\n",
    "                )\n",
    "            ]\n",
    "        )\n",
    "        .cpu()\n",
    "        .numpy()\n",
    "    )\n",
    "    if task_type == 'regression':\n",
    "        # Transform the predictions back to the original label space.\n",
    "        assert regression_label_stats is not None\n",
    "        y_pred = y_pred * regression_label_stats.std + regression_label_stats.mean\n",
    "\n",
    "    # Compute the mean of the k predictions.\n",
    "    if task_type != 'regression':\n",
    "        # For classification, the mean must be computed in the probabily space.\n",
    "        y_pred = scipy.special.softmax(y_pred, axis=-1)\n",
    "    y_pred = y_pred.mean(1)\n",
    "\n",
    "    y_true = data[part]['y'].cpu().numpy()\n",
    "    score = (\n",
    "        -(sklearn.metrics.mean_squared_error(y_true, y_pred) ** 0.5)\n",
    "        if task_type == 'regression'\n",
    "        else sklearn.metrics.accuracy_score(y_true, y_pred.argmax(1))\n",
    "    )\n",
    "    return float(score)  # The higher -- the better.\n",
    "\n",
    "\n",
    "print(f'Test score before training: {evaluate(\"test\"):.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 106.79it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.6059 (test) -0.6176\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 175.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5822 (test) -0.5930\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 296.53it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5604 (test) -0.5727\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 3: 100%|█████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 67.42it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5529 (test) -0.5626\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 4: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 131.65it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5446 (test) -0.5568\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 5: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 185.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5405 (test) -0.5485\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 6: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 250.92it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5294 (test) -0.5414\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 7: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 110.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5223 (test) -0.5290\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 8: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 166.18it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5199 (test) -0.5310\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 9: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 118.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5128 (test) -0.5218\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 10: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 276.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5126 (test) -0.5222\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 11: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 88.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5071 (test) -0.5143\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 12: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 208.86it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5120 (test) -0.5178\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 13: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 129.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5041 (test) -0.5131\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 14: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 290.99it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5035 (test) -0.5091\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 15: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 108.61it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5058 (test) -0.5122\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 16: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 192.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.5056 (test) -0.5109\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 17: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 172.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4974 (test) -0.5062\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 18: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 186.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4996 (test) -0.5092\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 19: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 205.61it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4931 (test) -0.4977\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 20: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 97.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4952 (test) -0.5010\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 21: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 212.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4935 (test) -0.4970\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 22: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 209.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4912 (test) -0.4964\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 23: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 287.73it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4885 (test) -0.4955\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 24: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 139.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4935 (test) -0.4979\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 25: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 139.68it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4929 (test) -0.4960\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 26: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 188.55it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4844 (test) -0.4889\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 27: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 198.41it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4857 (test) -0.4947\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 28: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 276.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4813 (test) -0.4902\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 29: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 189.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4885 (test) -0.4908\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 30: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 99.41it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4897 (test) -0.4973\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 31: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 113.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4819 (test) -0.4896\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 32: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 293.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4851 (test) -0.4920\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 33: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 226.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4832 (test) -0.4920\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 34: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 118.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4885 (test) -0.4930\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 35: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 164.63it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4848 (test) -0.4949\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 36: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 126.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4810 (test) -0.4885\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 37: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 295.54it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4834 (test) -0.4923\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 38: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 78.99it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4782 (test) -0.4861\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 39: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 178.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4792 (test) -0.4899\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 40: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 158.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4794 (test) -0.4907\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 41: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 271.73it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4819 (test) -0.4891\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 42: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 207.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4788 (test) -0.4864\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 43: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 84.36it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4790 (test) -0.4867\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 44: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 259.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4781 (test) -0.4894\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 45: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 119.45it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4866 (test) -0.4990\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 46: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 283.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4744 (test) -0.4826\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 47: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 112.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4766 (test) -0.4897\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 48: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 138.86it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4787 (test) -0.4861\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 49: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 139.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4743 (test) -0.4876\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 50: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 220.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4753 (test) -0.4855\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 51: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 241.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4828 (test) -0.5005\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 52: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 95.45it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4781 (test) -0.4925\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 53: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 204.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4705 (test) -0.4822\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 54: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 151.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4746 (test) -0.4845\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 55: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 285.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4746 (test) -0.4876\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 56: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 235.29it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4713 (test) -0.4879\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 57: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 128.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4707 (test) -0.4851\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 58: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 168.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4717 (test) -0.4861\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 59: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 97.40it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4714 (test) -0.4829\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 60: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 218.06it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4732 (test) -0.4869\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 61: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 73.24it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4729 (test) -0.4854\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 62: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 146.68it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4750 (test) -0.4867\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 63: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 225.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4713 (test) -0.4860\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 64: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 270.89it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4755 (test) -0.4840\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 65: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 70.83it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4741 (test) -0.4875\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 66: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 121.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4720 (test) -0.4877\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 67: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 207.24it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4699 (test) -0.4837\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 68: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 272.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4730 (test) -0.4899\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 69: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 168.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4690 (test) -0.4859\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 70: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 164.39it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4724 (test) -0.4865\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 71: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 100.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4704 (test) -0.4861\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 72: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 260.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4706 (test) -0.4827\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 73: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 250.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4726 (test) -0.4860\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 74: 100%|████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 90.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4722 (test) -0.4854\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 75: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 220.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4705 (test) -0.4838\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 76: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 174.36it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4709 (test) -0.4860\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 77: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 276.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4700 (test) -0.4844\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 78: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 106.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4682 (test) -0.4918\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 79: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 160.73it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4782 (test) -0.4930\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 80: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 159.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4697 (test) -0.4874\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 81: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 283.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4697 (test) -0.4842\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 82: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 156.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4749 (test) -0.4898\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 83: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 139.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4701 (test) -0.4902\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 84: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 157.91it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4684 (test) -0.4812\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 85: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 146.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4706 (test) -0.4885\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 86: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 253.91it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4700 (test) -0.4890\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 87: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 152.45it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4692 (test) -0.4864\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 88: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 161.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4747 (test) -0.4916\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 89: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 188.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4680 (test) -0.4833\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 90: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 229.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4673 (test) -0.4869\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 91: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 206.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4715 (test) -0.4916\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 92: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 109.96it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4737 (test) -0.4942\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 93: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 192.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4682 (test) -0.4849\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 94: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 133.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4705 (test) -0.4865\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 95: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 270.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4729 (test) -0.4889\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 96: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 238.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4698 (test) -0.4900\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 97: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 121.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4672 (test) -0.4878\n",
      "🌸 New best epoch! 🌸\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 98: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 167.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4706 (test) -0.4878\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 99: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 160.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4699 (test) -0.4846\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 100: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 240.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4779 (test) -0.4983\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 101: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 265.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4702 (test) -0.4889\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 102: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 115.55it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4693 (test) -0.4845\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 103: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 245.44it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4720 (test) -0.4917\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 104: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 113.64it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4692 (test) -0.4856\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 105: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 284.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4718 (test) -0.4908\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 106: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 212.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4721 (test) -0.4927\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 107: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 84.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4713 (test) -0.4905\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 108: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 276.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4691 (test) -0.4882\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 109: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 143.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4741 (test) -0.4924\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 110: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 284.04it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4711 (test) -0.4907\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 111: 100%|███████████████████████████████████████████████████████| 52/52 [00:00<00:00, 71.24it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4710 (test) -0.4866\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 112: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 169.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4700 (test) -0.4865\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 113: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 167.54it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4696 (test) -0.4878\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 114: 100%|██████████████████████████████████████████████████████| 52/52 [00:00<00:00, 255.78it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(val) -0.4728 (test) -0.4880\n",
      "\n",
      "\n",
      "Result:\n",
      "{'val': -0.4672268917866186, 'test': -0.4878477537379267, 'epoch': 97}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# For demonstration purposes (fast training and bad performance),\n",
    "# one can set smaller values:\n",
    "# n_epochs = 20\n",
    "# patience = 2\n",
    "n_epochs = 1_000_000_000\n",
    "patience = 16\n",
    "\n",
    "batch_size = 256\n",
    "epoch_size = math.ceil(len(train_idx) / batch_size)\n",
    "best = {\n",
    "    'val': -math.inf,\n",
    "    'test': -math.inf,\n",
    "    'epoch': -1,\n",
    "}\n",
    "# Early stopping: the training stops when\n",
    "# there are more than `patience` consequtive bad updates.\n",
    "patience = 16\n",
    "remaining_patience = patience\n",
    "\n",
    "print('-' * 88 + '\\n')\n",
    "for epoch in range(n_epochs):\n",
    "    for batch_idx in tqdm(\n",
    "        torch.randperm(len(data['train']['y']), device=device).split(batch_size),\n",
    "        desc=f'Epoch {epoch}',\n",
    "        total=epoch_size,\n",
    "    ):\n",
    "        model.train()\n",
    "        optimizer.zero_grad()\n",
    "        loss = loss_fn(apply_model('train', batch_idx), Y_train[batch_idx])\n",
    "        if grad_scaler is None:\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        else:\n",
    "            grad_scaler.scale(loss).backward()  # type: ignore\n",
    "            grad_scaler.step(optimizer)\n",
    "            grad_scaler.update()\n",
    "\n",
    "    val_score = evaluate('val')\n",
    "    test_score = evaluate('test')\n",
    "    print(f'(val) {val_score:.4f} (test) {test_score:.4f}')\n",
    "\n",
    "    if val_score > best['val']:\n",
    "        print('🌸 New best epoch! 🌸')\n",
    "        best = {'val': val_score, 'test': test_score, 'epoch': epoch}\n",
    "        remaining_patience = patience\n",
    "    else:\n",
    "        remaining_patience -= 1\n",
    "\n",
    "    if remaining_patience < 0:\n",
    "        break\n",
    "\n",
    "    print()\n",
    "\n",
    "print('\\n\\nResult:')\n",
    "print(best)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}