tutorials/training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example training script"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial shows a brief example of how the networks were trained. The specific example below is for $\\Delta$-learning of formation energy in a single-task setting. We'll start with the imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import h5py\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from delfta.net import EGNN\n",
    "from delfta.net_utils import MODEL_HPARAMS\n",
    "from delfta.utils import DATA_PATH, ROOT_PATH\n",
    "from torch_geometric.data import Data, DataLoader, Dataset\n",
    "from torch_geometric.utils import add_self_loops\n",
    "from torch_geometric.utils.undirected import to_undirected\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll download the training data if that hasn't happened yet (not done by default during the setup of `delfta`, since the original training files aren't needed to run the trained models). This can take a bit of time..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(os.path.join(DATA_PATH, \"qmugs\", \"qmugs_conf00.h5\")): \n",
    "    os.makedirs(os.path.join(DATA_PATH), exist_ok=True)\n",
    "    from delfta.download import DATASET_REMOTE, download\n",
    "    import tarfile \n",
    "    \n",
    "    download(DATASET_REMOTE, os.path.join(DATA_PATH, \"qmugs.tar.gz\"))\n",
    "    with tarfile.open(os.path.join(DATA_PATH, \"qmugs.tar.gz\")) as handle:\n",
    "        handle.extractall(DATA_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll define a new dataset class which is similar to the one in `delfta.net_utils`, but it doesn't load all the molecules in memory - which is better for training with a large number of molecules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DatasetSingletaskh5(Dataset):\n",
    "    def __init__(\n",
    "        self, txtfile, prop,\n",
    "    ):\n",
    "\n",
    "        # read txt\n",
    "        with open(txtfile, \"r\") as f:\n",
    "            chembls = [elem.rstrip(\"\\n\") for elem in f.readlines()]\n",
    "\n",
    "        # create dict on the fly: idx -> chembl\n",
    "        nums = list(range(0, len(chembls)))\n",
    "        self.idx2chembl = {}\n",
    "        for x in range(len(chembls)):\n",
    "            dict = {nums[x]: chembls[x]}\n",
    "            self.idx2chembl.update(dict)\n",
    "\n",
    "        # read h5\n",
    "        self.h5f0 = h5py.File(os.path.join(DATA_PATH, \"qmugs\", \"qmugs_conf00.h5\"), \"r\")\n",
    "        self.h5f1 = h5py.File(os.path.join(DATA_PATH, \"qmugs\", \"qmugs_conf01.h5\"), \"r\")\n",
    "        self.h5f2 = h5py.File(os.path.join(DATA_PATH, \"qmugs\", \"qmugs_conf02.h5\"), \"r\")\n",
    "\n",
    "        # define property of interest\n",
    "        self.prop = prop\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "\n",
    "        chembl_id = self.idx2chembl[idx]\n",
    "\n",
    "        #### nodes coordinates and target\n",
    "        if \"conf_00\" in chembl_id:\n",
    "            atomids = torch.LongTensor(self.h5f0[str(chembl_id)][\"atomids\"])\n",
    "            coords = torch.FloatTensor(self.h5f0[str(chembl_id)][\"coords\"])\n",
    "            target = torch.FloatTensor(self.h5f0[str(chembl_id)][self.prop])\n",
    "        elif \"conf_01\" in chembl_id:\n",
    "            atomids = torch.LongTensor(self.h5f1[str(chembl_id)][\"atomids\"])\n",
    "            coords = torch.FloatTensor(self.h5f1[str(chembl_id)][\"coords\"])\n",
    "            target = torch.FloatTensor(self.h5f1[str(chembl_id)][self.prop])\n",
    "        elif \"conf_02\" in chembl_id:\n",
    "            atomids = torch.LongTensor(self.h5f2[str(chembl_id)][\"atomids\"])\n",
    "            coords = torch.FloatTensor(self.h5f2[str(chembl_id)][\"coords\"])\n",
    "            target = torch.FloatTensor(self.h5f2[str(chembl_id)][self.prop])\n",
    "\n",
    "        #### edges\n",
    "        edge_index = np.array(nx.complete_graph(atomids.size(0)).edges())\n",
    "        edge_index = to_undirected(torch.from_numpy(edge_index).t().contiguous())\n",
    "        edge_index, _ = add_self_loops(edge_index, num_nodes=coords.shape[0])\n",
    "\n",
    "        #### graph object\n",
    "        graph_data = Data(\n",
    "            atomids=atomids,\n",
    "            coords=coords,\n",
    "            edge_index=edge_index,\n",
    "            target=target,\n",
    "            num_nodes=atomids.size(0),\n",
    "        )\n",
    "\n",
    "        return graph_data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.idx2chembl)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll define the training and evaluation loops:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "mae_loss = lambda x, y: F.l1_loss(x, y).item()\n",
    "\n",
    "\n",
    "def train_loop(model, loader, optimizer, criterion):\n",
    "    model.train()\n",
    "    training_loss = []\n",
    "\n",
    "    for g_batch in tqdm(loader, total=len(loader)):\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        g_batch = g_batch.to(DEVICE)\n",
    "        target = g_batch.target\n",
    "\n",
    "        prediction = model(g_batch).squeeze(1)\n",
    "\n",
    "        loss = criterion(prediction, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        with torch.no_grad():\n",
    "            mae = mae_loss(prediction, target)\n",
    "            training_loss.append(mae)\n",
    "\n",
    "    return np.mean(training_loss), np.std(training_loss)\n",
    "\n",
    "\n",
    "def eval_loop(model, loader):\n",
    "    model.eval()\n",
    "\n",
    "    maes = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for g_batch in tqdm(loader, total=len(loader)):\n",
    "            g_batch = g_batch.to(DEVICE)\n",
    "            prediction = model(g_batch).squeeze(1)\n",
    "            maes.append(mae_loss(prediction, g_batch.target))\n",
    "\n",
    "    eval_mae = sum(maes) / len(maes)\n",
    "    return eval_mae"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally we can start training the model:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50/50 [00:03<00:00, 14.05it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.06it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 0: 0.705199\n",
      "Epoch 2/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.86it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.62it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.21it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.36it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 2: 0.692088\n",
      "Epoch 4/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.94it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.56it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 3: 0.423953\n",
      "Epoch 5/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.63it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.14it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.61it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 5: 0.361086\n",
      "Epoch 7/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.00it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 6: 0.359345\n",
      "Epoch 8/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.95it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.34it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.14it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 8: 0.300412\n",
      "Epoch 10/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.00it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 11/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.05it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 12/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.37it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.11it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 13/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.20it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 14/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.00it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 15/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.72it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.57it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 16/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.18it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 17/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.96it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 18/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.79it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.85it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 19/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.16it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 20/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.02it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 21/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.82it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 20: 0.245532\n",
      "Epoch 22/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.83it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.31it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 23/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.09it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 24/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.24it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.57it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 25/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.89it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 26/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.30it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.56it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 27/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.11it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 28/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.17it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 29/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.26it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 30/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.11it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 31/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.00it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.64it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 32/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.88it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.58it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 31: 0.220015\n",
      "Epoch 33/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.88it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.17it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 34/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.05it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 35/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.14it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 36/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.27it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 37/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.97it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.56it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 38/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.23it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.77it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 39/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.25it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.53it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 40/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.65it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 41/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.78it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.97it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 42/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.05it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.88it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 43/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.08it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 44/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.24it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.54it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 45/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.77it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.83it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 46/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.89it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 47/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.02it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.55it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 48/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.77it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 49/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.15it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.47it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 50/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.96it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 51/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.99it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.93it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 52/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.83it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 53/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.18it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.08it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch    53: reducing learning rate of group 0 to 7.0000e-05.\n",
      "Epoch 54/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.12it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 55/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.02it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 56/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.19it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 57/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.11it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.58it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 58/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.94it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 59/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.11it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.11it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 60/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.26it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.65it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 61/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.02it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.14it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 62/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.05it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 63/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.02it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.16it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 64/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.94it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 65/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.16it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.47it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 66/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 16.24it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 67/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.82it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 68/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.85it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.81it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 67: 0.217697\n",
      "Epoch 69/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.85it/s]\n",
      "100%|██████████| 50/50 [00:02<00:00, 22.81it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 70/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 14.34it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.67it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 71/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.50it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.06it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 72/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 14.36it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 73/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 14.80it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 72: 0.211806\n",
      "Epoch 74/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 14.82it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.58it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 75/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.38it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 76/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 12.54it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 26.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 77/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.08it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 78/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.20it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 79/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 13.79it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.28it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 80/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.23it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 81/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.19it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 25.98it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 82/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.11it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 83/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.05it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 26.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 82: 0.186776\n",
      "Epoch 84/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 14.81it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 85/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 14.73it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 86/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.14it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.51it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 87/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 14.86it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.11it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 88/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.03it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 89/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 14.93it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.04it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 90/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.35it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 26.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 91/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.16it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 92/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 13.71it/s]\n",
      "100%|██████████| 50/50 [00:02<00:00, 24.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 93/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:04<00:00, 12.44it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.07it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 94/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.37it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 29.06it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New min eval_mae in epoch 93: 0.176719\n",
      "Epoch 95/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.15it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.06it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 96/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.51it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.64it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 97/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.24it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 98/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.11it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 28.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 99/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.17it/s]\n",
      "100%|██████████| 50/50 [00:01<00:00, 27.76it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 100/100...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 50/50 [00:03<00:00, 15.77it/s]\n",
      "100%|██████████| 50/50 [00:02<00:00, 24.31it/s]\n"
     ]
    }
   ],
   "source": [
    "save_path = os.path.join(ROOT_PATH, \"tutorials\", \"training_evaluation\")\n",
    "os.makedirs(save_path, exist_ok=True)\n",
    "train_path = os.path.join(DATA_PATH, \"qmugs\", \"train_example.txt\") # this includes only 50 molecules so that in this tutorial, we quickly see results\n",
    "eval_path = os.path.join(DATA_PATH, \"qmugs\", \"eval_example.txt\")  \n",
    "prop = \"DELTA_ENERGY\"\n",
    "model_param = MODEL_HPARAMS[\"single_energy_delta\"]\n",
    "\n",
    "train_data = DatasetSingletaskh5(txtfile=train_path, prop=prop)\n",
    "train_loader = DataLoader(train_data, batch_size=2, shuffle=True, num_workers=0)\n",
    "\n",
    "validation_data = DatasetSingletaskh5(txtfile=eval_path, prop=prop)\n",
    "validation_loader = DataLoader(\n",
    "    validation_data, batch_size=2, shuffle=False, num_workers=0\n",
    ")\n",
    "\n",
    "# Load model\n",
    "model = EGNN(\n",
    "    n_outputs=model_param.n_outputs,\n",
    "    global_prop=model_param.global_prop,\n",
    "    n_kernels=model_param.n_kernels,\n",
    "    mlp_dim=model_param.mlp_dim,\n",
    ")\n",
    "model = model.to(DEVICE)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-10)\n",
    "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n",
    "    optimizer, mode=\"min\", factor=0.7, patience=20, verbose=True\n",
    ")\n",
    "\n",
    "train_m_losses, train_std_losses = [], []\n",
    "val_losses = []\n",
    "val_maes = []\n",
    "epoch_maes = []\n",
    "min_mae = 1e10\n",
    "EPOCHS = 100\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    print(f\"Epoch {epoch + 1}/{EPOCHS}...\", flush=True)\n",
    "\n",
    "    m_loss, std_loss = train_loop(model, train_loader, optimizer, nn.MSELoss())\n",
    "    train_m_losses.append(m_loss)\n",
    "    train_std_losses.append(std_loss)\n",
    "\n",
    "    eval_mae = eval_loop(model, validation_loader)\n",
    "    val_losses.append(eval_mae)\n",
    "    scheduler.step(eval_mae)\n",
    "\n",
    "    if eval_mae < min_mae:\n",
    "\n",
    "        min_mae = eval_mae\n",
    "        val_maes.append(eval_mae)\n",
    "        epoch_maes.append(epoch)\n",
    "        print(f\"New min eval_mae in epoch {epoch}: {eval_mae:.6f}\", flush=True)\n",
    "        torch.save(\n",
    "            model.state_dict(), os.path.join(save_path, \"model.pt\"),\n",
    "        )\n",
    "        torch.save(\n",
    "            [train_m_losses, train_std_losses, val_losses],\n",
    "            os.path.join(save_path, \"loss_train_eval.pt\"),\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "6e3f0e0686e4e7c699d1571d1296d5e4355a8e109c70abfba87c453ed20075c6"
  },
  "kernelspec": {
   "display_name": "Python 3.7.4 64-bit ('delfta': conda)",
   "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.7.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}