testing_mask_edge_and_node.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import math\n",
    "import gc\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import torch\n",
    "from torch_geometric.data import Data\n",
    "\n",
    "tqdm.pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4bc0fab5be094712bc511a4fba6dafc1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_samples = 1000\n",
    "min_nodes = 4\n",
    "max_nodes = 20\n",
    "min_edges = 4\n",
    "\n",
    "data_list = []\n",
    "\n",
    "for i in tqdm(range(n_samples)):\n",
    "    n_nodes = random.randint(min_nodes, max_nodes)\n",
    "    x_features = torch.randn((n_nodes, 3))\n",
    "    x_mask = torch.zeros((n_nodes, 1))\n",
    "    x = torch.cat([x_features, x_mask], dim=1)\n",
    "        \n",
    "#     max_edges = math.factorial(n_nodes)\n",
    "    n_edges = random.randint(min_edges, n_nodes * 2)\n",
    "    edge_index = torch.randint(0, n_nodes - 1, (2, n_edges)).long()\n",
    "    \n",
    "    direction = torch.tensor([[int(edge_index[0][i] > edge_index[1][i])] for i in range(n_edges)])\n",
    "    \n",
    "    edge_attr = torch.cat([torch.randn((n_edges, 1)), direction, torch.zeros((n_edges, 2))], dim=1)\n",
    "\n",
    "    data = Data(\n",
    "        x=x,\n",
    "        edge_index=edge_index,\n",
    "        edge_attr=edge_attr\n",
    "    )\n",
    "    \n",
    "    data_list.append(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_list[0].x.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_list[0].edge_attr.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "CONFIG = {\n",
    "    \"device\": 0,\n",
    "    \"batch_size\": 256,\n",
    "    \"epochs\": 100,\n",
    "    \"lr\": 0.001,\n",
    "    \"decay\": 0,\n",
    "    \"num_layer\": 5,\n",
    "    \"emb_dim\": 300,\n",
    "    \"dropout_ratio\": 0,\n",
    "    \"mask_rate\": 0.15,\n",
    "    \"mask_edge\": 1,\n",
    "    \"JK\": \"last\",\n",
    "    \"output_model_file\": '',\n",
    "    \"gnn_type\": \"gat\",\n",
    "    \"seed\": 0, \n",
    "    \"num_workers\": 8\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "CONFIG[\"num_layer\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_accuracy(pred, target):\n",
    "    # return float(torch.sum((pred.detach() > 0) == target.to(torch.uint8)).cpu().item())/(pred.shape[0]*pred.shape[1])\n",
    "    return float(torch.sum(torch.max(pred.detach(), dim=1)[1] == target).cpu().item()) / len(pred)\n",
    "\n",
    "\n",
    "def train(config, model_list, loader, optimizer_list, device):\n",
    "    model, linear_pred_nodes, linear_pred_edges = model_list\n",
    "    optimizer_model, optimizer_linear_pred_nodes, optimizer_linear_pred_edges = optimizer_list\n",
    "\n",
    "    model.train()\n",
    "    linear_pred_nodes.train()\n",
    "    linear_pred_edges.train()\n",
    "\n",
    "    loss_accum = 0\n",
    "    n_broken_batches = 0\n",
    "    \n",
    "    pbar = tqdm(loader, desc=\"Iteration\")\n",
    "    for step, batch in enumerate(pbar):\n",
    "        batch = batch.to(device)\n",
    "        \n",
    "        optimizer_model.zero_grad()\n",
    "        optimizer_linear_pred_nodes.zero_grad()\n",
    "        optimizer_linear_pred_edges.zero_grad()\n",
    "                \n",
    "        node_rep = model(batch.x, batch.edge_index, batch.edge_attr)\n",
    "\n",
    "        ## loss for nodes\n",
    "        pred_node = linear_pred_nodes(node_rep[batch.masked_atom_indices])\n",
    "        node_loss = criterion(pred_node.double()[:, :3].double(), batch.mask_node_label[:, :3].double()) # FIX_HARDCODE: change [:, 0] to [:, edge_features]\n",
    "        # ADD: Computation of CrossEntropy for direction besides MSE\n",
    "\n",
    "#         if config[\"mask_edge\"]:\n",
    "        masked_edge_index = batch.edge_index[:, batch.connected_edge_indices]\n",
    "        edge_rep = node_rep[masked_edge_index[0]] + node_rep[masked_edge_index[1]]\n",
    "        pred_edge = linear_pred_edges(edge_rep)\n",
    "        edge_loss = criterion(pred_edge[:, 0].double(), batch.mask_edge_label[:, 0].double())\n",
    "        \n",
    "        \n",
    "        loss = node_loss + edge_loss\n",
    "        clip_loss = loss if not torch.isnan(loss).item() else torch.tensor(0.)\n",
    "        \n",
    "        if torch.isnan(clip_loss).item():\n",
    "            pbar.set_description(f\"Loss broken at batch {step}\")\n",
    "            \n",
    "        if not torch.isnan(loss).item():\n",
    "            clip_loss.backward() \n",
    "        \n",
    "            \n",
    "        optimizer_model.step()\n",
    "        optimizer_linear_pred_nodes.step()\n",
    "        optimizer_linear_pred_edges.step()\n",
    "\n",
    "        loss_accum += float(clip_loss.cpu().item())\n",
    "\n",
    "    return loss_accum / (step - n_broken_batches + 1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num layer: 5 mask rate: 0.150000 mask edge: 1\n",
      "====epoch 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:09<00:00, 101.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.336591654641272\n",
      "====epoch 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:10<00:00, 99.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.4817203775001981\n",
      "====epoch 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:10<00:00, 99.89it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0812932890001419\n",
      "====epoch 4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:11<00:00, 83.64it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.7920759572587341\n",
      "====epoch 5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:12<00:00, 80.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.584521589180429\n",
      "====epoch 6\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:10<00:00, 96.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.5135306873167448\n",
      "====epoch 7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:11<00:00, 89.69it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.43756665352361784\n",
      "====epoch 8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:11<00:00, 84.21it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.3094179191416715\n",
      "====epoch 9\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:12<00:00, 80.46it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.20913092897078733\n",
      "====epoch 10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:10<00:00, 92.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.17572478171183362\n",
      "====epoch 11\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:10<00:00, 93.68it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.15396351133588615\n",
      "====epoch 12\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:13<00:00, 74.63it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.1001390416327362\n",
      "====epoch 13\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:11<00:00, 88.44it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.07609405718440633\n",
      "====epoch 14\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:13<00:00, 72.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.06984401855973632\n",
      "====epoch 15\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:10<00:00, 91.73it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.05247504480403931\n",
      "====epoch 16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:12<00:00, 82.19it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.049419813835804155\n",
      "====epoch 17\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:12<00:00, 82.42it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.037639664728007935\n",
      "====epoch 18\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:12<00:00, 77.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.02651876103607247\n",
      "====epoch 19\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:11<00:00, 83.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.017810189705753595\n",
      "====epoch 20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:10<00:00, 90.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.020373847197847912\n",
      "====epoch 21\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:11<00:00, 85.58it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.021239763640060687\n",
      "====epoch 22\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:11<00:00, 84.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.01449511127969365\n",
      "====epoch 23\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:12<00:00, 80.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.006238851035727358\n",
      "====epoch 24\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:11<00:00, 87.24it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.007865678989143256\n",
      "====epoch 25\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:12<00:00, 81.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.008520478007851745\n",
      "====epoch 26\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:12<00:00, 81.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.006325016595180712\n",
      "====epoch 27\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:12<00:00, 79.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.004960495994406378\n",
      "====epoch 28\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:13<00:00, 74.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0018593676583555162\n",
      "====epoch 29\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:10<00:00, 93.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.003824493461943\n",
      "====epoch 30\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:11<00:00, 84.49it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0026027713690228607\n",
      "====epoch 31\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:13<00:00, 73.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0020737236056797906\n",
      "====epoch 32\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:12<00:00, 79.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.004386730244406535\n",
      "====epoch 33\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:15<00:00, 63.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0025653951540240195\n",
      "====epoch 34\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:15<00:00, 63.78it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0006536351096098714\n",
      "====epoch 35\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration:  37%|████████████████████████████████████████████▎                                                                            | 366/1000 [00:06<00:08, 74.30it/s]"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "\n",
    "from chem.dataloader import DataLoaderMasking  # , DataListLoader\n",
    "from chem.mydataset import MyDataset\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "from chem.model import GNN, GNN_graphpred\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from chem.util import MaskNode\n",
    "\n",
    "from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool\n",
    "\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "import timeit\n",
    "import warnings \n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "device = torch.device(\"cuda:\" + str(CONFIG[\"device\"])) if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(0)\n",
    "\n",
    "print(\"num layer: %d mask rate: %f mask edge: %d\" %(CONFIG[\"num_layer\"], CONFIG[\"mask_rate\"], CONFIG[\"mask_edge\"]))\n",
    "\n",
    "dataset = MyDataset(data_list, transform=MaskNode(num_node_features=data_list[0].x.shape[1],\n",
    "                                                  num_edge_features=data_list[0].edge_attr.shape[1],\n",
    "                                                  mask_rate=CONFIG[\"mask_rate\"], \n",
    "                                                  mask_edge=CONFIG[\"mask_edge\"]))\n",
    "\n",
    "loader = DataLoaderMasking(dataset, batch_size=1, shuffle=True, num_workers=0)\n",
    "\n",
    "model = GNN(CONFIG[\"num_layer\"], \n",
    "            CONFIG[\"emb_dim\"], \n",
    "            x_input_dim = data_list[0].x.shape[1],\n",
    "            edge_attr_input_dim = data_list[0].edge_attr.shape[1],\n",
    "            JK = CONFIG[\"JK\"], \n",
    "            drop_ratio = CONFIG[\"dropout_ratio\"], \n",
    "            gnn_type = CONFIG[\"gnn_type\"]).to(device)\n",
    "linear_pred_nodes = torch.nn.Linear(CONFIG[\"emb_dim\"], data_list[0].x.shape[1]).to(device)\n",
    "linear_pred_edges = torch.nn.Linear(CONFIG[\"emb_dim\"], data_list[0].edge_attr.shape[1] - 2).to(device)\n",
    "\n",
    "model_list = [model, linear_pred_nodes, linear_pred_edges]\n",
    "\n",
    "# set up optimizers\n",
    "optimizer_model = optim.Adam(model.parameters(), lr=CONFIG[\"lr\"], weight_decay=CONFIG[\"decay\"])\n",
    "optimizer_linear_pred_nodes = optim.Adam(linear_pred_nodes.parameters(), lr=CONFIG[\"lr\"], weight_decay=CONFIG[\"decay\"])\n",
    "optimizer_linear_pred_edges = optim.Adam(linear_pred_edges.parameters(), lr=CONFIG[\"lr\"], weight_decay=CONFIG[\"decay\"])\n",
    "\n",
    "optimizer_list = [optimizer_model, optimizer_linear_pred_nodes, optimizer_linear_pred_edges]\n",
    "\n",
    "for epoch in range(1, CONFIG[\"epochs\"] + 1):\n",
    "    print(\"====epoch \" + str(epoch))\n",
    "\n",
    "    train_loss = train(CONFIG, model_list, loader, optimizer_list, device)\n",
    "    print(train_loss)\n",
    "\n",
    "if not CONFIG[\"model_file\"] == \"\":\n",
    "    torch.save(model.state_dict(), CONFIG[\"model_file\"] + \".pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'chem.batch' from '/Users/user/PycharmProjects/pretrain_gnns/chem/batch.py'>"
      ]
     },
     "execution_count": 162,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import importlib\n",
    "from chem import model, util, batch\n",
    "\n",
    "importlib.reload(model)\n",
    "importlib.reload(util)\n",
    "importlib.reload(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat \"chem/batch.py\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}