project_solubility_final_transfer_learning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "from tdc.single_pred.adme import ADME\n",
    "from tdc import Evaluator\n",
    "from tqdm.notebook import tqdm, trange\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem\n",
    "from matplotlib import pyplot as plt\n",
    "from IPython import display\n",
    "\n",
    "from typing import List, Tuple\n",
    "\n",
    "\n",
    "class Featurizer:\n",
    "    def __init__(self, y_column, smiles_col='Drug', **kwargs):\n",
    "        self.y_column = y_column\n",
    "        self.smiles_col = smiles_col\n",
    "        self.__dict__.update(kwargs)\n",
    "    \n",
    "    def __call__(self, df):\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.data import Data\n",
    "from torch_geometric.loader import DataLoader\n",
    "\n",
    "\n",
    "def one_of_k_encoding(x, allowable_set):\n",
    "    if x not in allowable_set:\n",
    "        raise ValueError(\"input {0} not in allowable set{1}:\".format(\n",
    "            x, allowable_set))\n",
    "    return list(map(lambda s: x == s, allowable_set))\n",
    "\n",
    "\n",
    "def one_of_k_encoding_unk(x, allowable_set):\n",
    "    if x not in allowable_set:\n",
    "        x = allowable_set[-1]\n",
    "    return list(map(lambda s: x == s, allowable_set))\n",
    "\n",
    "\n",
    "from rdkit.Chem import rdMolDescriptors\n",
    "\n",
    "class GraphFeaturizer(Featurizer):\n",
    "    def __call__(self, df, getRepresentation):\n",
    "        graphs = []\n",
    "        labels = []\n",
    "        for i, row in df.iterrows():\n",
    "            y = row[self.y_column]\n",
    "            smiles = row[self.smiles_col]\n",
    "            mol = Chem.MolFromSmiles(smiles)\n",
    "            \n",
    "            edges = []\n",
    "            for bond in mol.GetBonds():\n",
    "                begin = bond.GetBeginAtomIdx()\n",
    "                end = bond.GetEndAtomIdx()\n",
    "                edges.append((begin, end))  # TODO: Add edges in both directions\n",
    "            edges = np.array(edges)\n",
    "            \n",
    "            nodes = []\n",
    "            for atom in mol.GetAtoms():\n",
    "                # print(atom.GetAtomicNum(), atom.GetNumImplicitHs(), atom.GetTotalNumHs(), atom.GetSymbol(), atom.GetNumExplicitHs(), atom.GetTotalValence())\n",
    "                results = getRepresentation(atom)\n",
    "                # print(results)\n",
    "                nodes.append(results)\n",
    "            nodes = np.array(nodes)\n",
    "            \n",
    "            graphs.append((nodes, edges.T))\n",
    "            labels.append(y)\n",
    "        labels = np.array(labels)\n",
    "        return [Data(\n",
    "            x=torch.FloatTensor(x), \n",
    "            edge_index=torch.LongTensor(edge_index), \n",
    "            y=torch.FloatTensor([y])\n",
    "        ) for ((x, edge_index), y) in zip(graphs, labels)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def defaultRepresentation(atom):\n",
    "    return one_of_k_encoding_unk(atom.GetAtomicNum(), range(11)) + one_of_k_encoding(\n",
    "                    atom.GetDegree(), range(11)\n",
    "                ) + one_of_k_encoding_unk(\n",
    "                    atom.GetImplicitValence(), range(11)\n",
    "                ) + [atom.GetIsAromatic()] + one_of_k_encoding_unk(\n",
    "                    atom.GetTotalNumHs(), range(11)\n",
    "                ) + [atom.GetNumImplicitHs(), atom.GetFormalCharge(), atom.GetNumRadicalElectrons(), atom.IsInRing()] # TODO: Add atom features as a list, you can use one_of_k_encodings defined above\n",
    "\n",
    "def representation1(atom):\n",
    "    return one_of_k_encoding_unk(atom.GetAtomicNum(), range(12)) + one_of_k_encoding_unk(\n",
    "                    atom.GetDegree(), range(6)) + one_of_k_encoding_unk(\n",
    "                    atom.GetTotalNumHs(), range(5)\n",
    "                ) + [atom.GetFormalCharge(), atom.IsInRing(), atom.GetIsAromatic()]\n",
    "\n",
    "def representation10(atom):\n",
    "    return one_of_k_encoding_unk(atom.GetAtomicNum(), range(12)) + one_of_k_encoding_unk(\n",
    "                    atom.GetDegree(), range(6)) + one_of_k_encoding_unk(\n",
    "                    atom.GetTotalNumHs(), range(5)\n",
    "                ) + [atom.IsInRing(), atom.GetIsAromatic()]\n",
    "\n",
    "def representationAll(atom):\n",
    "    return one_of_k_encoding_unk(atom.GetAtomicNum(), range(12)) + one_of_k_encoding_unk(\n",
    "                    atom.GetDegree(), range(6)) + one_of_k_encoding_unk(\n",
    "                    atom.GetTotalNumHs(), range(5)) + one_of_k_encoding_unk(\n",
    "                    atom.GetImplicitValence(), range(6))  + one_of_k_encoding_unk(\n",
    "                    atom.GetHybridization(),\n",
    "                    [\n",
    "                        Chem.rdchem.HybridizationType.SP, Chem.rdchem.HybridizationType.SP2,\n",
    "                        Chem.rdchem.HybridizationType.SP3, Chem.rdchem.HybridizationType.SP3D,\n",
    "                        Chem.rdchem.HybridizationType.SP3D2\n",
    "                    ]\n",
    "                ) + [atom.GetFormalCharge(), atom.IsInRing(), atom.GetIsAromatic()\n",
    "                     ] + [atom.GetNumRadicalElectrons()]\n",
    "\n",
    "def printProperties(atom):\n",
    "    print(\"=========\")\n",
    "    print(\"GetDegree\", atom.GetDegree())\n",
    "    print(\"GetImplicitValence\", atom.GetImplicitValence())\n",
    "    print(\"GetAtomicNum\", atom.GetAtomicNum())\n",
    "    print(\"GetTotalNumHs\", atom.GetTotalNumHs())\n",
    "    print(\"GetNumImplicitHs\", atom.GetNumImplicitHs())\n",
    "    print(\"GetNeighbors\", atom.GetNeighbors())\n",
    "    print(\"GetNumExplicitHs\", atom.GetNumExplicitHs())\n",
    "    print(\"GetTotalDegree\", atom.GetTotalDegree())\n",
    "    print(\"GetTotalNumHs\", atom.GetTotalNumHs())\n",
    "    print(\"GetTotalValence\", atom.GetTotalValence())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found local copy...\n",
      "Loading...\n",
      "Done!\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:53] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:54] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:59] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:59] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:07:59] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:00] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:00] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:00] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:01] WARNING: not removing hydrogen atom without neighbors\n"
     ]
    }
   ],
   "source": [
    "class ECFPFeaturizer(Featurizer):\n",
    "    def __init__(self, y_column, radius=2, length=1024, **kwargs):\n",
    "        self.radius = radius\n",
    "        self.length = length\n",
    "        super().__init__(y_column, **kwargs)\n",
    "    \n",
    "    def __call__(self, df):\n",
    "        fingerprints = []\n",
    "        labels = []\n",
    "        for i, row in df.iterrows():\n",
    "            y = row[self.y_column]\n",
    "            smiles = row[self.smiles_col]\n",
    "            mol = Chem.MolFromSmiles(smiles)\n",
    "            fp = AllChem.GetMorganFingerprintAsBitVect(mol, self.radius, nBits=self.length)\n",
    "            fingerprints.append(fp)\n",
    "            labels.append(y)\n",
    "        fingerprints = np.array(fingerprints)\n",
    "        labels = np.array(labels)\n",
    "        return fingerprints, labels\n",
    "\n",
    "data = ADME('Solubility_AqSolDB')\n",
    "split = data.get_split()\n",
    "rmse = Evaluator(name = 'RMSE')\n",
    "\n",
    "featurizer = ECFPFeaturizer(y_column='Y')\n",
    "X_train, y_train = featurizer(split['train'])\n",
    "X_valid, y_valid = featurizer(split['valid'])\n",
    "X_test, y_test = featurizer(split['test'])\n",
    "\n",
    "featurizer = GraphFeaturizer('Y')\n",
    "graph = featurizer(split['test'].iloc[:1], defaultRepresentation)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:03] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:04] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:05] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:05] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:05] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:05] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:05] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:05] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:06] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:06] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:08] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:08] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:08] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:08] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:08] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:08] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:09] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:10] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:10] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:10] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:10] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:10] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:10] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:10] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:10] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:10] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:11] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:12] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:12] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:12] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:12] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:12] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:12] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:12] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:12] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:12] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:12] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:14] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:15] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:15] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:15] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:15] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:15] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:15] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:16] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:17] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:18] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:19] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:19] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:19] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:19] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:19] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:19] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:20] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:20] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:22] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:22] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:22] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:22] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:23] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:23] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:23] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:23] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:23] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n",
      "[15:08:24] WARNING: not removing hydrogen atom without neighbors\n"
     ]
    }
   ],
   "source": [
    "from torch_geometric.loader import DataLoader as GraphDataLoader\n",
    "\n",
    "\n",
    "# prepare data loaders\n",
    "batch_size = 64\n",
    "\n",
    "#dla repr1\n",
    "train_loader1 = GraphDataLoader(featurizer(split['train'], representation1), batch_size=batch_size, shuffle=True)\n",
    "valid_loader1 = GraphDataLoader(featurizer(split['valid'], representation1), batch_size=batch_size)\n",
    "test_loader1 = GraphDataLoader(featurizer(split['test'], representation1), batch_size=batch_size)\n",
    "\n",
    "#dla repr10\n",
    "train_loader10 = GraphDataLoader(featurizer(split['train'], representation10), batch_size=batch_size, shuffle=True)\n",
    "valid_loader10 = GraphDataLoader(featurizer(split['valid'], representation10), batch_size=batch_size)\n",
    "test_loader10 = GraphDataLoader(featurizer(split['test'], representation10), batch_size=batch_size)\n",
    "\n",
    "#dla naszej warstwy\n",
    "train_loader = GraphDataLoader(featurizer(split['train'], representationAll), batch_size=batch_size, shuffle=True)\n",
    "valid_loader = GraphDataLoader(featurizer(split['valid'], representationAll), batch_size=batch_size)\n",
    "test_loader = GraphDataLoader(featurizer(split['test'], representationAll), batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GCNConv, GINConv, global_mean_pool\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import global_mean_pool as gap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#warstwa attention pooling\n",
    "class MyAttentionModule3(torch.nn.Module): # zakladamy ze atom ma 49 featerow\n",
    "    def __init__(self, groupFeatures=1):\n",
    "        super().__init__()\n",
    "        self.groupFeatures = groupFeatures\n",
    "        self.gates = torch.nn.ModuleDict({ # do wyliczenia atencji dla kazdej grupy cech - jest ich 9\n",
    "            'AtomicNum': GCNConv(12, 1),\n",
    "            'Degree': GCNConv(6, 1),\n",
    "            'TotalNumHs': GCNConv(5, 1),\n",
    "            'ImplicitValence': GCNConv(6, 1),\n",
    "            'Hybridization': GCNConv(5, 1),\n",
    "            'FormalCharge': GCNConv(1, 1),\n",
    "            'IsInRing': GCNConv(1, 1),\n",
    "            'IsAromatic': GCNConv(1, 1),\n",
    "            'NumRadicalElectrons': GCNConv(1, 1)\n",
    "        })\n",
    "        \n",
    "        self.feats = torch.nn.ModuleDict({ # do transformacji grupy cech w wektor, na razie dziala tylko dla groupFeatures=1\n",
    "            'AtomicNum': torch.nn.Linear(12, groupFeatures),\n",
    "            'Degree': torch.nn.Linear(6, groupFeatures),\n",
    "            'TotalNumHs': torch.nn.Linear(5, groupFeatures),\n",
    "            'ImplicitValence': torch.nn.Linear(6, groupFeatures),\n",
    "            'Hybridization': torch.nn.Linear(5, groupFeatures),\n",
    "            'FormalCharge': torch.nn.Linear(1, groupFeatures),\n",
    "            'IsInRing': torch.nn.Linear(1, groupFeatures),\n",
    "            'IsAromatic': torch.nn.Linear(1, groupFeatures),\n",
    "            'NumRadicalElectrons': torch.nn.Linear(1, groupFeatures)\n",
    "        })\n",
    "\n",
    "    def forward(self, x, edge_index, batch):\n",
    "        gates = []\n",
    "        gates.append(self.gates['AtomicNum'](x[:,0:12], edge_index))\n",
    "        gates.append(self.gates['Degree'](x[:,12:18], edge_index))\n",
    "        gates.append(self.gates['TotalNumHs'](x[:,18:23], edge_index))\n",
    "        gates.append(self.gates['ImplicitValence'](x[:,23:29], edge_index))\n",
    "        gates.append(self.gates['Hybridization'](x[:,29:34], edge_index))\n",
    "        gates.append(self.gates['FormalCharge'](x[:,34:35], edge_index))\n",
    "        gates.append(self.gates['IsInRing'](x[:,35:36], edge_index))\n",
    "        gates.append(self.gates['IsAromatic'](x[:,36:37], edge_index))\n",
    "        gates.append(self.gates['NumRadicalElectrons'](x[:,37:38], edge_index))\n",
    "        logits = torch.cat(gates, dim=-1)\n",
    "        attention = torch.softmax(logits, dim=-1).unsqueeze(-1)\n",
    "        \n",
    "        subgroups = []\n",
    "        subgroups.append(self.feats['AtomicNum'](x[:,0:12]) * attention[:,0])\n",
    "        subgroups.append(self.feats['Degree'](x[:,12:18]) * attention[:,1])\n",
    "        subgroups.append(self.feats['TotalNumHs'](x[:,18:23]) * attention[:,2])\n",
    "        subgroups.append(self.feats['ImplicitValence'](x[:,23:29]) * attention[:,3])\n",
    "        subgroups.append(self.feats['Hybridization'](x[:,29:34]) * attention[:,4])\n",
    "        subgroups.append(self.feats['FormalCharge'](x[:,34:35]) * attention[:,5])\n",
    "        subgroups.append(self.feats['IsInRing'](x[:,35:36]) * attention[:,6])\n",
    "        subgroups.append(self.feats['IsAromatic'](x[:,36:37]) * attention[:,7])\n",
    "        subgroups.append(self.feats['NumRadicalElectrons'](x[:,37:38]) * attention[:,8])\n",
    "        x = torch.stack(subgroups, dim=-2)\n",
    "        x = torch.sum(x, dim=-2)\n",
    "        \n",
    "        return x, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def attSequential(n_feats):\n",
    "    return torch.nn.Sequential(torch.nn.Linear(n_feats, 1),\n",
    "                       torch.nn.BatchNorm1d(1), torch.nn.ReLU(),\n",
    "                       torch.nn.Linear(1, 1), torch.nn.ReLU())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#warstwa attention pooling\n",
    "class MyAttentionModule4(torch.nn.Module): # zakladamy ze atom ma 49 featerow\n",
    "    def __init__(self, groupFeatures=1):\n",
    "        super().__init__()\n",
    "        self.groupFeatures = groupFeatures\n",
    "        self.gates = torch.nn.ModuleDict({ # do wyliczenia atencji dla kazdej grupy cech - jest ich 9\n",
    "            'AtomicNum': GINConv(attSequential(12), train_eps=True),\n",
    "            'Degree': GINConv(attSequential(6), train_eps=True),\n",
    "            'TotalNumHs': GINConv(attSequential(5), train_eps=True),\n",
    "            'ImplicitValence': GINConv(attSequential(6), train_eps=True),\n",
    "            'Hybridization': GINConv(attSequential(5), train_eps=True),\n",
    "            'FormalCharge': GINConv(attSequential(1), train_eps=True),\n",
    "            'IsInRing': GINConv(attSequential(1), train_eps=True),\n",
    "            'IsAromatic': GINConv(attSequential(1), train_eps=True),\n",
    "            'NumRadicalElectrons': GINConv(attSequential(1), train_eps=True)\n",
    "        })\n",
    "        \n",
    "        self.feats = torch.nn.ModuleDict({ # do transformacji grupy cech w wektor, na razie dziala tylko dla groupFeatures=1\n",
    "            'AtomicNum': torch.nn.Linear(12, groupFeatures),\n",
    "            'Degree': torch.nn.Linear(6, groupFeatures),\n",
    "            'TotalNumHs': torch.nn.Linear(5, groupFeatures),\n",
    "            'ImplicitValence': torch.nn.Linear(6, groupFeatures),\n",
    "            'Hybridization': torch.nn.Linear(5, groupFeatures),\n",
    "            'FormalCharge': torch.nn.Linear(1, groupFeatures),\n",
    "            'IsInRing': torch.nn.Linear(1, groupFeatures),\n",
    "            'IsAromatic': torch.nn.Linear(1, groupFeatures),\n",
    "            'NumRadicalElectrons': torch.nn.Linear(1, groupFeatures)\n",
    "        })\n",
    "\n",
    "    def forward(self, x, edge_index, batch):\n",
    "        gates = []\n",
    "        gates.append(self.gates['AtomicNum'](x[:,0:12], edge_index))\n",
    "        gates.append(self.gates['Degree'](x[:,12:18], edge_index))\n",
    "        gates.append(self.gates['TotalNumHs'](x[:,18:23], edge_index))\n",
    "        gates.append(self.gates['ImplicitValence'](x[:,23:29], edge_index))\n",
    "        gates.append(self.gates['Hybridization'](x[:,29:34], edge_index))\n",
    "        gates.append(self.gates['FormalCharge'](x[:,34:35], edge_index))\n",
    "        gates.append(self.gates['IsInRing'](x[:,35:36], edge_index))\n",
    "        gates.append(self.gates['IsAromatic'](x[:,36:37], edge_index))\n",
    "        gates.append(self.gates['NumRadicalElectrons'](x[:,37:38], edge_index))\n",
    "        logits = torch.cat(gates, dim=-1)\n",
    "        attention = torch.softmax(logits, dim=-1).unsqueeze(-1)\n",
    "        \n",
    "        subgroups = []\n",
    "        subgroups.append(self.feats['AtomicNum'](x[:,0:12]) * attention[:,0])\n",
    "        subgroups.append(self.feats['Degree'](x[:,12:18]) * attention[:,1])\n",
    "        subgroups.append(self.feats['TotalNumHs'](x[:,18:23]) * attention[:,2])\n",
    "        subgroups.append(self.feats['ImplicitValence'](x[:,23:29]) * attention[:,3])\n",
    "        subgroups.append(self.feats['Hybridization'](x[:,29:34]) * attention[:,4])\n",
    "        subgroups.append(self.feats['FormalCharge'](x[:,34:35]) * attention[:,5])\n",
    "        subgroups.append(self.feats['IsInRing'](x[:,35:36]) * attention[:,6])\n",
    "        subgroups.append(self.feats['IsAromatic'](x[:,36:37]) * attention[:,7])\n",
    "        subgroups.append(self.feats['NumRadicalElectrons'](x[:,37:38]) * attention[:,8])\n",
    "        x = torch.stack(subgroups, dim=-2)\n",
    "        x = torch.sum(x, dim=-2)\n",
    "        \n",
    "        return x, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphNeuralNetwork(torch.nn.Module):  # TODO: assign hyperparameters to attributes and define the forward pass\n",
    "    def __init__(self, hidden_size, n_convs=3, my_layer=None, features_after_layer=26, n_features=49, dropout=0.2):\n",
    "        super().__init__()\n",
    "        self.myAttentionModule = my_layer\n",
    "        self.dropout = dropout\n",
    "\n",
    "        convs = torch.nn.ModuleList()\n",
    "        convs.append(GCNConv(features_after_layer, hidden_size))\n",
    "        for i in range(1, n_convs):\n",
    "            convs.append(GCNConv(hidden_size, hidden_size))\n",
    "        self.convs = convs\n",
    "        self.linear = torch.nn.Linear(hidden_size, 1)\n",
    "    \n",
    "    def forward(self, x, edge_index, batch):\n",
    "        att = None\n",
    "        if self.myAttentionModule is not None:\n",
    "            x, att = self.myAttentionModule(x, edge_index, batch)\n",
    "        for i in range(0, len(self.convs)-1):\n",
    "            x = self.convs[i](x, edge_index)\n",
    "            x = x.relu()\n",
    "        x = self.convs[-1](x, edge_index)\n",
    "        \n",
    "        x = gap(x, batch)\n",
    "        \n",
    "        x = F.dropout(x, p=self.dropout, training=self.training)\n",
    "\n",
    "        out = self.linear(x)\n",
    "\n",
    "        return out, att"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader, valid_loader, epochs=20, learning_rate = 0.01):\n",
    "    model.train()\n",
    "    \n",
    "    # training loop\n",
    "    optimizer = torch.optim.Adam(model.parameters(), learning_rate) # TODO: define an optimizer\n",
    "    loss_fn = torch.nn.MSELoss()  # TODO: define a loss function\n",
    "    for epoch in trange(1, epochs + 1, leave=False):\n",
    "        for data in tqdm(train_loader, leave=False):\n",
    "            x, edge_index, batch, y = data.x, data.edge_index, data.batch, data.y\n",
    "            model.zero_grad()\n",
    "            preds, att = model(x, edge_index, batch)\n",
    "            loss = loss_fn(preds, y.reshape(-1, 1))\n",
    "            loss.backward()\n",
    "            # print(\"==============\")\n",
    "            # for par in model.myAttentionModule.parameters():\n",
    "            #     print(par)\n",
    "            optimizer.step()\n",
    "    return model\n",
    "\n",
    "\n",
    "def predict(model, test_loader):\n",
    "    # evaluation loop\n",
    "    preds_batches = []\n",
    "    with torch.no_grad():\n",
    "        for data in tqdm(test_loader):\n",
    "            x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "            \n",
    "            preds, att = model(x, edge_index, batch)\n",
    "            preds_batches.append(preds.cpu().detach().numpy())\n",
    "    preds = np.concatenate(preds_batches)\n",
    "    return preds, att"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "\n",
    "def train_best(model, train_loader, valid_loader, epochs=20, learning_rate = 0.01, saveImg=False, title=\"\"):\n",
    "    model.train()\n",
    "\n",
    "    best_state = deepcopy(model.state_dict())\n",
    "    best_val = 1000000\n",
    "    \n",
    "    # training loop\n",
    "    optimizer = torch.optim.Adam(model.parameters(), learning_rate) # TODO: define an optimizer\n",
    "    loss_fn = torch.nn.MSELoss()  # TODO: define a loss function\n",
    "    for epoch in trange(1, epochs + 1, leave=False):\n",
    "        # preds_batches = []\n",
    "        running_loss = 0.0\n",
    "        for data in train_loader:\n",
    "            x, edge_index, batch, y = data.x, data.edge_index, data.batch, data.y\n",
    "            model.zero_grad()\n",
    "            preds, att = model(x, edge_index, batch)\n",
    "            loss = loss_fn(preds, y.reshape(-1, 1))\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        # evaluation loop\n",
    "        preds_batches = []\n",
    "        with torch.no_grad():\n",
    "            for data in valid_loader:\n",
    "                x, edge_index, batch, y = data.x, data.edge_index, data.batch, data.y\n",
    "                preds, att = model(x, edge_index, batch)\n",
    "                loss = loss_fn(preds, y.reshape(-1, 1))\n",
    "                preds_batches.append(preds.cpu().detach().numpy())\n",
    "        preds = np.concatenate(preds_batches)\n",
    "        mae = rmse(y_valid, preds.flatten())\n",
    "        if mae < best_val:\n",
    "            best_state = deepcopy(model.state_dict())\n",
    "            best_val = mae\n",
    "            print(best_val)\n",
    "\n",
    "    model.load_state_dict(best_state)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize(model, train_loader, valid_loader, test_loader, epochs=20, learning_rate = 0.01, saveImg=False, title=\"\"):\n",
    "    model.train()\n",
    "\n",
    "    torch.save(model, \"train.pth\")\n",
    "    best_val = 1000000\n",
    "    \n",
    "    # training loop\n",
    "    optimizer = torch.optim.Adam(model.parameters(), learning_rate) # TODO: define an optimizer\n",
    "    loss_fn = torch.nn.MSELoss()  # TODO: define a loss function\n",
    "    train_losses = []\n",
    "    val_losses = []\n",
    "    train_errors = []\n",
    "    val_errors = []\n",
    "    for epoch in trange(1, epochs + 1, leave=False):\n",
    "        # preds_batches = []\n",
    "        running_loss = 0.0\n",
    "        for data in train_loader:\n",
    "            x, edge_index, batch, y = data.x, data.edge_index, data.batch, data.y\n",
    "            model.zero_grad()\n",
    "            preds, att = model(x, edge_index, batch)\n",
    "            loss = loss_fn(preds, y.reshape(-1, 1))\n",
    "            # print(len(train_dataset))\n",
    "\n",
    "            running_loss += loss.item()\n",
    "            # preds_batches.append(preds.cpu().detach().numpy())\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        epoch_loss = running_loss / len(train_loader)\n",
    "        train_losses.append(epoch_loss)\n",
    "        # preds = np.concatenate(preds_batches)\n",
    "        # mae = rmse(y_train, preds.flatten())\n",
    "        # train_errors.append(mae)\n",
    "\n",
    "        # evaluation loop\n",
    "        preds_batches = []\n",
    "        running_loss = 0.0\n",
    "        with torch.no_grad():\n",
    "            for data in valid_loader:\n",
    "                x, edge_index, batch, y = data.x, data.edge_index, data.batch, data.y\n",
    "                preds, att = model(x, edge_index, batch)\n",
    "                loss = loss_fn(preds, y.reshape(-1, 1))\n",
    "                # print(len(train_dataset))\n",
    "\n",
    "                running_loss += loss.item()\n",
    "                preds_batches.append(preds.cpu().detach().numpy())\n",
    "        epoch_loss = running_loss / len(valid_loader)\n",
    "        val_losses.append(epoch_loss)\n",
    "        preds = np.concatenate(preds_batches)\n",
    "        mae = rmse(y_valid, preds.flatten())\n",
    "        if mae < best_val:\n",
    "            torch.save(model, \"train.pth\")\n",
    "            best_val = mae\n",
    "            print(best_val)\n",
    "        val_errors.append(mae)\n",
    "\n",
    "    model = torch.load(\"train.pth\")\n",
    "    model.eval()\n",
    "\n",
    "    ##### visualize ########\n",
    "    plt.plot(train_losses, label='train_loss')\n",
    "    plt.plot(val_losses, label='val_loss')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    if saveImg:\n",
    "        plt.savefig(title + \"_loss.png\")\n",
    "\n",
    "    # plt.plot(train_errors,label='train_errors')\n",
    "    plt.plot(val_errors, label='val_RMSE')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    if saveImg:\n",
    "        plt.savefig(title + \"_val_error.png\")\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame({\"Repr 1\": [], \"Repr 10\": [],\n",
    "                   \"transfer learning - size = 3\": [], \"transfer learning - size = 35\": [], \"transfer learning - size = 100\": [], \"transfer learning - size = 35 (big)\": []})\n",
    "pd.set_option(\"display.precision\", 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cdf9f06f435140ed898d5b0d17d5f442",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.9239051927224744\n",
      "1.8824044650857024\n",
      "1.8708188823695409\n",
      "1.8544374099999812\n",
      "1.8514084692328474\n",
      "1.8505132278028962\n",
      "1.8445121074406807\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "80890fd0baf6465bbb3392a0cf168eba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d345673852274435a9989600c97812f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.9202650505641574\n",
      "1.899382670135747\n",
      "1.8779640098929185\n",
      "1.8764937760804201\n",
      "1.8734432236981806\n",
      "1.8459894030322437\n",
      "1.8440723721801575\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77e69915ebb84514bda78fd6f7707793",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4c6df0bc27fd445c81b4606bf5291114",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.2111081134628714\n",
      "2.1993403799033966\n",
      "2.198601756081361\n",
      "2.196553355901469\n",
      "2.191896473762494\n",
      "2.188236792649506\n",
      "2.1867627697859713\n",
      "2.184694445927528\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "25b4c6d2ce9d4b6ca4c8ceb8ba499dd8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2e6e4f1cd2a0409a891c510bd4deed67",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.006275511074213\n",
      "1.9686244163744882\n",
      "1.9251508140100526\n",
      "1.923495775654457\n",
      "1.9055809401847608\n",
      "1.8875762879330271\n",
      "1.883742406027809\n",
      "1.883458562269263\n",
      "1.8710887540855654\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "68b49b00633d45ada79919921a6ef590",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3d7c54134e4d44299e858f33959f1fa3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0883162272902998\n",
      "2.026305533327386\n",
      "1.9498733497591907\n",
      "1.9203216628288267\n",
      "1.916607533693913\n",
      "1.891485258106759\n",
      "1.8524325079111086\n",
      "1.8392726819319016\n",
      "1.8325868467407926\n",
      "1.826979008309253\n",
      "1.8259048853623259\n",
      "1.8103511658036004\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "210eaaa494334e3698105ec6b551a385",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "240a7c4084b84225955efd2782c18c8b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8689479907031568\n",
      "1.8490449941565912\n",
      "1.8432180988148414\n",
      "1.8319006974358878\n",
      "1.8237140759264738\n",
      "1.821234907804118\n",
      "1.8206668762305838\n",
      "1.8080845297523491\n",
      "1.801659534220267\n",
      "1.7993488850243853\n",
      "1.7959151707860204\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a51db156d3b2404287b8ba0f5d16328b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d0513ab9b8e845fa95706f091f6b5e71",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8988952916050632\n",
      "1.8854509774977821\n",
      "1.850625755708985\n",
      "1.8365761139196664\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "35a55df4a23b451298c340a4212cd797",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "00e7caa3b067436fa72da03b8e93a11f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.9090209076090874\n",
      "1.8502024077249433\n",
      "1.8382101219101776\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "270a31581490438fa71504df1d5bca09",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "65de0bf2f65548d1841a4da4485939ae",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.2093757941969154\n",
      "2.1913067330062987\n",
      "2.1875101048045082\n",
      "2.1815202434928493\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5dfb5cf18daa4fc49315447b45e6d4aa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d74656ab7baf42caa4bd3d4843221784",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.970370867360132\n",
      "1.950321733365964\n",
      "1.9314371930781502\n",
      "1.9242463064462412\n",
      "1.91024940179593\n",
      "1.9011162618334152\n",
      "1.8805715890360317\n",
      "1.87994124342532\n",
      "1.8755903687203994\n",
      "1.8727136330028433\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "05637fa7ba7445019206c81df5a5e939",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b382b9ccc2d6467b88ffff137882a00e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.6280118934304175\n",
      "3.970491568325811\n",
      "2.2687763661780735\n",
      "2.1948151260074464\n",
      "2.162550729127905\n",
      "2.06865168470698\n",
      "2.05740753144236\n",
      "1.9460267186743494\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0ec305b5577f4a5abcac0a1eaab1bf08",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d0906fb428d4ee88f1455a3cf7f663f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.9325431860851146\n",
      "1.8653189326206043\n",
      "1.844039909564728\n",
      "1.83749829831167\n",
      "1.8336068994056027\n",
      "1.8233891175558552\n",
      "1.8043976619234836\n",
      "1.8032968224597137\n",
      "1.789958707763793\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "97a1a5207bf6433e89301de423ddf9b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6abf9324eae44ca4b177abf62b9c1f5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8135539298562566\n",
      "1.6993360878639885\n",
      "1.6885876697918238\n",
      "1.601957954766087\n",
      "1.5858535669793297\n",
      "1.5645577292669195\n",
      "1.5502527651410325\n",
      "1.5428647062993208\n",
      "1.5211438253774008\n",
      "1.5202397624620274\n",
      "1.5067218538558627\n",
      "1.4849105412565649\n",
      "1.4603146348714047\n",
      "1.4566700959288572\n",
      "1.4246726958126805\n",
      "1.4244575319883075\n",
      "1.4104772663651108\n",
      "1.4026701419119965\n",
      "1.3984617187106438\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2ff5a22dda5240a2a950fcaca925c459",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b34f81e09ef04edb8001d7703d7687d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8870590871032598\n",
      "1.7450174320543508\n",
      "1.6783174575772293\n",
      "1.6255218454440337\n",
      "1.5553863122463154\n",
      "1.5326127397416378\n",
      "1.4915281034202208\n",
      "1.4886664975425312\n",
      "1.487519923635393\n",
      "1.4577398251955764\n",
      "1.4501861137896384\n",
      "1.426146831878994\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "31162f53c3214865ba4b6a42b4ee5e24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ed9e2e723344217888433f39f591e9b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.4562389586878655\n",
      "2.0401166836517257\n",
      "1.9692780801315766\n",
      "1.9138323444742749\n",
      "1.89181664030575\n",
      "1.8617221888077689\n",
      "1.835579179253183\n",
      "1.8043014216946491\n",
      "1.7823326193501703\n",
      "1.7431856030289545\n",
      "1.7422829361715828\n",
      "1.7138130360074986\n",
      "1.7055674275233266\n",
      "1.6961986024054339\n",
      "1.6717113428028303\n",
      "1.6659618119448998\n",
      "1.6463086137953833\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c130dba875ee47fdbb9bcbf5f695efa2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f4ad21408fcd4e8a97c4711e3fca6df1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.9426676686155255\n",
      "1.75723625561017\n",
      "1.7198558775396031\n",
      "1.694935809797026\n",
      "1.6876341525739398\n",
      "1.5993223362028088\n",
      "1.593086700769341\n",
      "1.5752532193197215\n",
      "1.5597870882572984\n",
      "1.545270637221411\n",
      "1.5431074660433017\n",
      "1.5031868879115222\n",
      "1.494317095980469\n",
      "1.4900655417052573\n",
      "1.4873994935700299\n",
      "1.4794641153552224\n",
      "1.467616902129646\n",
      "1.4671442647662178\n",
      "1.4605251442989984\n",
      "1.442238944487148\n",
      "1.4412246612600559\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e6d9ea8e82b44b47b0245e81161e6622",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ecefc69764924817817d6bb720566fd0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.9711632287931635\n",
      "1.84641310424506\n",
      "1.7343743059991\n",
      "1.721062344735718\n",
      "1.68118665363238\n",
      "1.6534378300037589\n",
      "1.6219635998079953\n",
      "1.5994969227695182\n",
      "1.5820379921100873\n",
      "1.5729723834465144\n",
      "1.5467106443421113\n",
      "1.5152639646234154\n",
      "1.5128128982045634\n",
      "1.4888630112262697\n",
      "1.4883967018727462\n",
      "1.4725466723903282\n",
      "1.45521455362106\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "559e3eff1a37416795bb28d5cb8b0a85",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bb910fd7aeda48848369a5e5e9c8b3d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.7192836897023978\n",
      "1.6883662289845778\n",
      "1.6112009138366723\n",
      "1.6012756760824105\n",
      "1.5214404698565986\n",
      "1.5091839785209995\n",
      "1.4918499506992586\n",
      "1.4705950291591636\n",
      "1.4449996549252466\n",
      "1.4422145933262327\n",
      "1.4097203751545624\n",
      "1.3849929879586094\n",
      "1.382792694785252\n",
      "1.3757132042996427\n",
      "1.3686928932617943\n",
      "1.3656276844200645\n",
      "1.3642610852322599\n",
      "1.3472988044361898\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d3a51154311f424d88a713ffdb30b802",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "380155d6c85244158e6037fd245c804d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.7386515185411178\n",
      "1.7122691873648843\n",
      "1.6972600994321485\n",
      "1.5669464182517663\n",
      "1.5646584156267176\n",
      "1.5438415417358486\n",
      "1.5102544345902689\n",
      "1.5022343525916215\n",
      "1.4821937003799173\n",
      "1.4717586172538215\n",
      "1.461083535010673\n",
      "1.4443587139822627\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e3ddb4da05d042c08ecacb7bddef9cbf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "93c1a4e96985434bb8bda860cd04b840",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0484212651355653\n",
      "1.7624432019545655\n",
      "1.7359618541963509\n",
      "1.6498225788741978\n",
      "1.6090717352451496\n",
      "1.5923067535190116\n",
      "1.5544581858539133\n",
      "1.5439842005401705\n",
      "1.539299858740487\n",
      "1.515136254660383\n",
      "1.5107798426199908\n",
      "1.5024722687763312\n",
      "1.4804395006807345\n",
      "1.460315694969541\n",
      "1.4545631838505713\n",
      "1.4351245881840857\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f66223231884719985c596fad94a67e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb6e68f8375647dfb8b6da221d12cf73",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.1856532298948004\n",
      "2.0872137795617425\n",
      "1.9504013919724328\n",
      "1.9128047235374106\n",
      "1.9126929836073416\n",
      "1.9052292084965141\n",
      "1.838107377835481\n",
      "1.7896863407532195\n",
      "1.7771637431716132\n",
      "1.7622401775521694\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5649b247f4214e0faf28d871f159d0a1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "557fcd99e11449a29aa6ae52f6f060ad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.7512769930377114\n",
      "1.7264509575966038\n",
      "1.6395988432686879\n",
      "1.6242661585748468\n",
      "1.5549478482881267\n",
      "1.5438167142500043\n",
      "1.4984894657795043\n",
      "1.4977881104636748\n",
      "1.4487460241275996\n",
      "1.4417491720927849\n",
      "1.4300594732494316\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9bb7e0f7250d473c8db763cdaf5e25d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0a5467c2121b470f932a88e3f4406112",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8362439321075437\n",
      "1.8047134412541495\n",
      "1.7532059623888383\n",
      "1.7050720349688016\n",
      "1.6929390577521106\n",
      "1.679922206440848\n",
      "1.6766625161810473\n",
      "1.6690362549816442\n",
      "1.6686289807975552\n",
      "1.602742267641253\n",
      "1.5972219808822494\n",
      "1.564652769705757\n",
      "1.5638882923233122\n",
      "1.560677552483208\n",
      "1.5588284579819658\n",
      "1.5244305177436648\n",
      "1.519306788958806\n",
      "1.5085519327061232\n",
      "1.4982027579227692\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e24c60a1af974221a9c5e03095079d35",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "347b708c086543c880c447508101aeee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.75629808089122\n",
      "1.6486641637361596\n",
      "1.6236028173821315\n",
      "1.5526594292479599\n",
      "1.4963945314887694\n",
      "1.4393138267512864\n",
      "1.4149987039725263\n",
      "1.4140297400874158\n",
      "1.378447142298283\n",
      "1.3642425588599543\n",
      "1.355407947932431\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e8c3a700e54423eb7b19155f8e3f4da",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ecd528693b02462687dae357a816b467",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.7496708435074497\n",
      "1.720067526633254\n",
      "1.6758601563571933\n",
      "1.6426997636729983\n",
      "1.6087502440827193\n",
      "1.5529745324415023\n",
      "1.5330511885246785\n",
      "1.5190697452905537\n",
      "1.498257140019345\n",
      "1.4833248546241218\n",
      "1.478494095050053\n",
      "1.4556834153910758\n",
      "1.4456846409869661\n",
      "1.4130737330824672\n",
      "1.3943385144468634\n",
      "1.3942283505941149\n",
      "1.3921790647771597\n",
      "1.3841391911848444\n",
      "1.3791118448337563\n",
      "1.3726702435473803\n",
      "1.345278514990747\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a3e8b44460104953b9c0f5e61875e428",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4a9844321206479a85329d69b84051f7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.841473083239969\n",
      "1.6784425315229696\n",
      "1.653437287827916\n",
      "1.6017856873674035\n",
      "1.5857702729788092\n",
      "1.5616181991692804\n",
      "1.5473871340321381\n",
      "1.544154978552965\n",
      "1.5312865205094544\n",
      "1.5290383479239238\n",
      "1.5156955591219567\n",
      "1.501473121203392\n",
      "1.4920576059603066\n",
      "1.4822912690265044\n",
      "1.4622476152630046\n",
      "1.428134086249175\n",
      "1.427308864194191\n",
      "1.4231815358690867\n",
      "1.4146146698997386\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fa7f84bc5625435b83114a33d788a60d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ed26d752214641e7b339ebf6d232e737",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.076069829012615\n",
      "1.9422019365821994\n",
      "1.9061508731731378\n",
      "1.865284115390599\n",
      "1.8105948962783305\n",
      "1.7254081775627776\n",
      "1.7196977171844017\n",
      "1.683280356324051\n",
      "1.648829552395884\n",
      "1.628170997782659\n",
      "1.5780329752313798\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6b30099baee34e66a0e1848c20a9e2ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1edc057133d441718cbb1cf00d9628fe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8168484098948792\n",
      "1.69816949199348\n",
      "1.6579548958102102\n",
      "1.5648560684426942\n",
      "1.5640926564483202\n",
      "1.5082131830561782\n",
      "1.500638363944748\n",
      "1.4977893917299003\n",
      "1.492077989938716\n",
      "1.4831611156923323\n",
      "1.455540419316024\n",
      "1.454935630780983\n",
      "1.4405598649285138\n",
      "1.4378425222455435\n",
      "1.4255530790854347\n",
      "1.413663133720055\n",
      "1.4034268933321985\n",
      "1.4010448027762437\n",
      "1.3922676125109394\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "41ebf518c4bf47e2b5a168086a557fc6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e9aa5763794a468ca226afd4393d3166",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.761623383504531\n",
      "1.8001863845251138\n",
      "1.7540693766505153\n",
      "1.733500637559276\n",
      "1.7058382213062966\n",
      "1.6430424410326319\n",
      "1.6210814107985392\n",
      "1.6013472732531626\n",
      "1.57334788993754\n",
      "1.5480634915641736\n",
      "1.5474392859383128\n",
      "1.5115309562214152\n",
      "1.510718536656948\n",
      "1.498976329114164\n",
      "1.4895567305638955\n",
      "1.4654136265153388\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a8ae84de170642ad9ba0f51db9a8ec1a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "42a9fa539ca24aada0938de37da91029",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8212545678417087\n",
      "1.716413478224713\n",
      "1.6160955660411274\n",
      "1.5693476552958594\n",
      "1.5577903579120422\n",
      "1.5265532732544573\n",
      "1.5227637738434905\n",
      "1.5108418994816195\n",
      "1.4426445488665347\n",
      "1.3972953867684548\n",
      "1.3885217456087402\n",
      "1.371540023212772\n",
      "1.3303096365320886\n",
      "1.322358792917672\n",
      "1.3212027699775033\n",
      "1.321119733441724\n",
      "1.315106577651445\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cc730e8902cf49b891e5c245a1538c13",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bf67e54ff7524909930cc45d391c4c89",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.7742731661431803\n",
      "1.7488160693516566\n",
      "1.6627376663102695\n",
      "1.6384430532143035\n",
      "1.6116838014219879\n",
      "1.5881892794801167\n",
      "1.5611366571218228\n",
      "1.5607767671289292\n",
      "1.5550249136722762\n",
      "1.5193065472484653\n",
      "1.5164415262288724\n",
      "1.4975248902732232\n",
      "1.4915679495352097\n",
      "1.4877241299701243\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3368b730e6334edbaa56700e3be5ac82",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "018d1249e7a3481fac3c11a0a6d5257a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.7854953530003375\n",
      "1.7637375888492972\n",
      "1.6173879909226458\n",
      "1.5839717056144904\n",
      "1.5590481876321147\n",
      "1.5387392867246168\n",
      "1.5378589196149515\n",
      "1.5348193394208363\n",
      "1.528878998686646\n",
      "1.5262852325114382\n",
      "1.5218132491412848\n",
      "1.5188419253915946\n",
      "1.502617151338612\n",
      "1.4995090606262158\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "24007bec86e643859104dfc1252f3376",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9625a6d29f9443e6901b01806886c172",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.324078738585975\n",
      "2.3049589978459086\n",
      "2.302687600117142\n",
      "2.3020991662540684\n",
      "2.3017426126310543\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9df052051bf547d6bba1ed907727cf94",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c2499baffd2f4f77b026e76a1c9c0135",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8781730843471336\n",
      "1.7799389993596189\n",
      "1.6901251063984473\n",
      "1.675658495170125\n",
      "1.580310039633441\n",
      "1.5136312680890247\n",
      "1.49626872824438\n",
      "1.4920869368197238\n",
      "1.4861852105687192\n",
      "1.4800051928296254\n",
      "1.4642672823049312\n",
      "1.4486449835224178\n",
      "1.421802838662332\n",
      "1.401142914591916\n",
      "1.378153882651602\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4abc40fba3e2498b9e8ff315722502f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "58028d5e918f45f7a8e4e51abf65d32b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.108788761067722\n",
      "2.009442597364585\n",
      "1.7914146399595603\n",
      "1.7868792554557367\n",
      "1.708368658133901\n",
      "1.65581426843549\n",
      "1.622608035668734\n",
      "1.6222855537463963\n",
      "1.5789159718380492\n",
      "1.5456348207998767\n",
      "1.5191277286633533\n",
      "1.499797238049101\n",
      "1.4945454973051917\n",
      "1.481274450458955\n",
      "1.4475974381485155\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bfeafe4538dd4f2b96f20bc9b8246da2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cbb0df9ef37d4092ac0f4c00da9c7161",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/70 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8292562102344363\n",
      "1.7136875793628301\n",
      "1.6168999905843453\n",
      "1.5911928172625052\n",
      "1.5781093224352765\n",
      "1.5566501737351732\n",
      "1.543287926304379\n",
      "1.5418125006934897\n",
      "1.5161119736026962\n",
      "1.5028243355653375\n",
      "1.4968223055806262\n",
      "1.4787897703191486\n",
      "1.4681521182603297\n",
      "1.4634721189720536\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f7732a7994d34f84b5a3692fe0220673",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_times = 1\n",
    "\n",
    "for n_convs in [1, 3, 5]:\n",
    "    for n_channels in [64, 512]:\n",
    "        row = []\n",
    "        #########################\n",
    "        scores = []\n",
    "        for _ in range(n_times):\n",
    "            m =  GraphNeuralNetwork(n_channels, n_convs=n_convs, features_after_layer=26)\n",
    "\n",
    "            m = train_best(m, train_loader1, valid_loader1, 70)\n",
    "            predictions, att = predict(m, test_loader1)\n",
    "            rmse_score = rmse(y_test, predictions.flatten())\n",
    "            scores.append(\"{:.2f}\".format(rmse_score))\n",
    "        row.append(\" | \".join(scores))\n",
    "\n",
    "        #########################\n",
    "        scores = []\n",
    "        for _ in range(n_times):\n",
    "            m =  GraphNeuralNetwork(n_channels, n_convs=n_convs, features_after_layer=25)\n",
    "\n",
    "            m = train_best(m, train_loader10, valid_loader10, 70)\n",
    "            predictions, att = predict(m, test_loader10)\n",
    "            rmse_score = rmse(y_test, predictions.flatten())\n",
    "            scores.append(\"{:.2f}\".format(rmse_score))\n",
    "        row.append(\" | \".join(scores))\n",
    "\n",
    "        #########################\n",
    "        for vect_size in [3, 35, 100]:\n",
    "            scores = []\n",
    "            for _ in range(n_times):\n",
    "                m =  GraphNeuralNetwork(n_channels, n_convs=n_convs, my_layer=MyAttentionModule4(vect_size), features_after_layer=vect_size)\n",
    "\n",
    "                m.myAttentionModule.load_state_dict(torch.load('attention_pooling' + str(vect_size) + '.pth'))\n",
    "                m.eval()\n",
    "                for par in m.myAttentionModule.parameters():\n",
    "                        par.requires_grad = False\n",
    "\n",
    "                m = train_best(m, train_loader, valid_loader, 70)\n",
    "                predictions, att = predict(m, test_loader)\n",
    "                rmse_score = rmse(y_test, predictions.flatten())\n",
    "                scores.append(\"{:.2f}\".format(rmse_score))\n",
    "            row.append(\" | \".join(scores))\n",
    "\n",
    "        #########################\n",
    "        scores = []\n",
    "        for _ in range(n_times):\n",
    "            m =  GraphNeuralNetwork(n_channels, n_convs=n_convs, my_layer=MyAttentionModule4(35), features_after_layer=35)\n",
    "\n",
    "            m.myAttentionModule.load_state_dict(torch.load('attention_pooling35_big.pth'))\n",
    "            m.eval()\n",
    "            for par in m.myAttentionModule.parameters():\n",
    "                    par.requires_grad = False\n",
    "\n",
    "            m = train_best(m, train_loader, valid_loader, 70)\n",
    "            predictions, att = predict(m, test_loader)\n",
    "            rmse_score = rmse(y_test, predictions.flatten())\n",
    "            scores.append(\"{:.2f}\".format(rmse_score))\n",
    "        row.append(\" | \".join(scores))\n",
    "\n",
    "        df.loc[str(n_convs) + \" convs, \" + str(n_channels) + \" channels\"] = row"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Repr 1</th>\n",
       "      <th>Repr 10</th>\n",
       "      <th>transfer learning - size = 3</th>\n",
       "      <th>transfer learning - size = 35</th>\n",
       "      <th>transfer learning - size = 100</th>\n",
       "      <th>transfer learning - size = 35 (big)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1 convs, 64 channels</th>\n",
       "      <td>1.90</td>\n",
       "      <td>1.90</td>\n",
       "      <td>2.22</td>\n",
       "      <td>1.92</td>\n",
       "      <td>1.95</td>\n",
       "      <td>1.79</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1 convs, 512 channels</th>\n",
       "      <td>1.92</td>\n",
       "      <td>1.90</td>\n",
       "      <td>2.24</td>\n",
       "      <td>1.93</td>\n",
       "      <td>2.17</td>\n",
       "      <td>1.78</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3 convs, 64 channels</th>\n",
       "      <td>1.40</td>\n",
       "      <td>1.46</td>\n",
       "      <td>1.65</td>\n",
       "      <td>1.42</td>\n",
       "      <td>1.47</td>\n",
       "      <td>1.33</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3 convs, 512 channels</th>\n",
       "      <td>1.43</td>\n",
       "      <td>1.43</td>\n",
       "      <td>1.75</td>\n",
       "      <td>1.41</td>\n",
       "      <td>1.50</td>\n",
       "      <td>1.35</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5 convs, 64 channels</th>\n",
       "      <td>1.39</td>\n",
       "      <td>1.43</td>\n",
       "      <td>1.60</td>\n",
       "      <td>1.36</td>\n",
       "      <td>1.47</td>\n",
       "      <td>1.67</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5 convs, 512 channels</th>\n",
       "      <td>1.50</td>\n",
       "      <td>1.56</td>\n",
       "      <td>2.33</td>\n",
       "      <td>1.38</td>\n",
       "      <td>1.50</td>\n",
       "      <td>1.46</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                      Repr 1 Repr 10 transfer learning - size = 3   \n",
       "1 convs, 64 channels    1.90    1.90                         2.22  \\\n",
       "1 convs, 512 channels   1.92    1.90                         2.24   \n",
       "3 convs, 64 channels    1.40    1.46                         1.65   \n",
       "3 convs, 512 channels   1.43    1.43                         1.75   \n",
       "5 convs, 64 channels    1.39    1.43                         1.60   \n",
       "5 convs, 512 channels   1.50    1.56                         2.33   \n",
       "\n",
       "                      transfer learning - size = 35   \n",
       "1 convs, 64 channels                           1.92  \\\n",
       "1 convs, 512 channels                          1.93   \n",
       "3 convs, 64 channels                           1.42   \n",
       "3 convs, 512 channels                          1.41   \n",
       "5 convs, 64 channels                           1.36   \n",
       "5 convs, 512 channels                          1.38   \n",
       "\n",
       "                      transfer learning - size = 100   \n",
       "1 convs, 64 channels                            1.95  \\\n",
       "1 convs, 512 channels                           2.17   \n",
       "3 convs, 64 channels                            1.47   \n",
       "3 convs, 512 channels                           1.50   \n",
       "5 convs, 64 channels                            1.47   \n",
       "5 convs, 512 channels                           1.50   \n",
       "\n",
       "                      transfer learning - size = 35 (big)  \n",
       "1 convs, 64 channels                                 1.79  \n",
       "1 convs, 512 channels                                1.78  \n",
       "3 convs, 64 channels                                 1.33  \n",
       "3 convs, 512 channels                                1.35  \n",
       "5 convs, 64 channels                                 1.67  \n",
       "5 convs, 512 channels                                1.46  "
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m =  GraphNeuralNetwork(n_channels, n_convs=n_convs, my_layer=MyAttentionModule4(vect_size), features_after_layer=vect_size)\n",
    "\n",
    "m.myAttentionModule.load_state_dict(torch.load('attention_pooling35_big.pth'))\n",
    "for par in m.myAttentionModule.parameters():\n",
    "        par.requires_grad = False\n",
    "\n",
    "m = train_best(m, train_loader, valid_loader, 70)\n",
    "predictions, att = predict(m, test_loader)\n",
    "rmse_score = rmse(y_test, predictions.flatten())\n",
    "###################### atencja #################################\n",
    "\n",
    "df_single = pd.DataFrame({\"AtomicNum\": [], \"Degree\": [], \"TotalNumHs\": [], \"ImplicitValence\": [], \"Hybridization\": [], \"FormalCharge\": [],\n",
    "                          \"IsInRing\": [], \"IsAromatic\": [], \"NumRadicalElectrons\": []})\n",
    "df_single.style.set_caption(\"Hello World\")\n",
    "\n",
    "df_batch = pd.DataFrame({\"AtomicNum\": [], \"Degree\": [], \"TotalNumHs\": [], \"ImplicitValence\": [], \"Hybridization\": [], \"FormalCharge\": [],\n",
    "                          \"IsInRing\": [], \"IsAromatic\": [], \"NumRadicalElectrons\": []})\n",
    "df_batch.style.set_caption(\"Hello World\")\n",
    "\n",
    "preds_batches = []\n",
    "with torch.no_grad():\n",
    "    for data in test_loader:\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        \n",
    "        preds, att = m(x, edge_index, batch)\n",
    "        preds_batches.append(preds.cpu().detach().numpy())\n",
    "        att = att.squeeze()\n",
    "        df_single.loc[len(df_single)] = att[0].tolist()\n",
    "        df_batch.loc[len(df_single)] = torch.mean(gap(att, batch), dim=0).tolist()\n",
    "preds = np.concatenate(preds_batches)\n",
    "\n",
    "rmse_score = rmse(y_test, predictions.flatten())\n",
    "\n",
    "print(f'RMSE = {rmse_score:.2f}')\n",
    "df_single[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_batch[:10]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mldd23",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}