notebooks/tutorial_hyper_keras_tuner.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a890e09c",
   "metadata": {},
   "source": [
    "# Hyperoptimization with keras-tuner"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44fc8ddc",
   "metadata": {},
   "source": [
    "## 1. Dataset\n",
    "\n",
    "Example for dataset ESOL from Moleculenet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d4f23fce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import keras as ks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "189277c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:kgcnn.data.download:Checking and possibly downloading dataset with name ESOL\n",
      "INFO:kgcnn.data.download:Dataset directory located at C:\\Users\\patri\\.kgcnn\\datasets\n",
      "INFO:kgcnn.data.download:Dataset directory found. Done.\n",
      "INFO:kgcnn.data.download:Dataset found. Done.\n",
      "INFO:kgcnn.data.ESOL:Found SDF C:\\Users\\patri\\.kgcnn\\datasets\\ESOL\\delaney-processed.sdf of pre-computed structures.\n",
      "INFO:kgcnn.data.ESOL:Read molecules from mol-file.\n",
      "INFO:kgcnn.data.ESOL: ... process molecules 0 from 1128\n",
      "INFO:kgcnn.molecule.encoder:OneHotEncoder Symbol found ['O', 'C', 'N', 'S', 'Cl', 'P', 'F', 'I', 'Br']\n",
      "INFO:kgcnn.molecule.encoder:OneHotEncoder Hybridization found [rdkit.Chem.rdchem.HybridizationType.SP3, rdkit.Chem.rdchem.HybridizationType.SP, rdkit.Chem.rdchem.HybridizationType.SP2]\n",
      "INFO:kgcnn.molecule.encoder:OneHotEncoder TotalDegree found [2, 4, 1, 3]\n",
      "INFO:kgcnn.molecule.encoder:OneHotEncoder TotalNumHs found [1, 2, 0, 3, 4]\n",
      "INFO:kgcnn.molecule.encoder:OneHotEncoder CIPCode found [None, 'S', 'R']\n",
      "INFO:kgcnn.molecule.encoder:OneHotEncoder ChiralityPossible found [None, '1']\n",
      "INFO:kgcnn.molecule.encoder:OneHotEncoder BondType found [rdkit.Chem.rdchem.BondType.SINGLE, rdkit.Chem.rdchem.BondType.TRIPLE, rdkit.Chem.rdchem.BondType.AROMATIC, rdkit.Chem.rdchem.BondType.DOUBLE]\n",
      "INFO:kgcnn.molecule.encoder:OneHotEncoder Stereo found [rdkit.Chem.rdchem.BondStereo.STEREONONE, rdkit.Chem.rdchem.BondStereo.STEREOE, rdkit.Chem.rdchem.BondStereo.STEREOZ]\n"
     ]
    }
   ],
   "source": [
    "from kgcnn.data.datasets.ESOLDataset import ESOLDataset\n",
    "dataset = ESOLDataset()\n",
    "# dataset.set_attributes()\n",
    "dataset.map_list(method=\"set_edge_indices_reverse\");\n",
    "dataset.map_list(method=\"count_nodes_and_edges\");\n",
    "dataset.map_list(**{\"method\": \"count_nodes_and_edges\", \"total_edges\": \"total_reverse\"});"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f537e620",
   "metadata": {},
   "source": [
    "## 2. Model\n",
    "\n",
    "Pick DMPNN as example model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f3cffd2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kgcnn.literature.DMPNN import make_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f34d2eb0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:kgcnn.models.utils:Updated model kwargs: '{'name': 'DMPNN', 'inputs': [{'shape': (None, 41), 'name': 'node_attributes', 'dtype': 'float32'}, {'shape': (None, 11), 'name': 'edge_attributes', 'dtype': 'float32'}, {'shape': (None, 2), 'name': 'edge_indices', 'dtype': 'int64'}, {'shape': (None, 1), 'name': 'edge_indices_reverse', 'dtype': 'int64'}, {'shape': (), 'name': 'total_nodes', 'dtype': 'int64'}, {'shape': (), 'name': 'total_edges', 'dtype': 'int64'}, {'shape': (), 'name': 'total_reverse', 'dtype': 'int64'}], 'input_tensor_type': 'padded', 'cast_disjoint_kwargs': {}, 'input_embedding': None, 'input_node_embedding': {'input_dim': 95, 'output_dim': 64}, 'input_edge_embedding': {'input_dim': 5, 'output_dim': 64}, 'input_graph_embedding': {'input_dim': 100, 'output_dim': 64}, 'pooling_args': {'pooling_method': 'scatter_sum'}, 'use_graph_state': False, 'edge_initialize': {'units': 128, 'use_bias': True, 'activation': 'relu'}, 'edge_dense': {'units': 128, 'use_bias': True, 'activation': 'linear'}, 'edge_activation': {'activation': 'relu'}, 'node_dense': {'units': 128, 'use_bias': True, 'activation': 'relu'}, 'verbose': 10, 'depth': 5, 'dropout': {'rate': 0.1}, 'output_embedding': 'graph', 'output_to_tensor': None, 'output_tensor_type': 'padded', 'output_mlp': {'use_bias': [True, True, False], 'units': [64, 32, 1], 'activation': ['relu', 'relu', 'linear']}, 'output_scaling': None}'.\n"
     ]
    }
   ],
   "source": [
    "model_config = {\n",
    "    \"name\": \"DMPNN\",\n",
    "    \"inputs\": [\n",
    "        {\"shape\": (None, 41), \"name\": \"node_attributes\", \"dtype\": \"float32\"},\n",
    "        {\"shape\": (None, 11), \"name\": \"edge_attributes\", \"dtype\": \"float32\"},\n",
    "        {\"shape\": (None, 2), \"name\": \"edge_indices\", \"dtype\": \"int64\"},\n",
    "        {\"shape\": (None, 1), \"name\": \"edge_indices_reverse\", \"dtype\": \"int64\"},\n",
    "        {\"shape\": (), \"name\": \"total_nodes\", \"dtype\": \"int64\"},\n",
    "        {\"shape\": (), \"name\": \"total_edges\", \"dtype\": \"int64\"},\n",
    "        {\"shape\": (), \"name\": \"total_reverse\", \"dtype\": \"int64\"}\n",
    "    ],\n",
    "    \"cast_disjoint_kwargs\": {},\n",
    "    \"input_node_embedding\": {\"input_dim\": 95, \"output_dim\": 64},\n",
    "    \"input_edge_embedding\": {\"input_dim\": 5, \"output_dim\": 64},\n",
    "    \"input_graph_embedding\": {\"input_dim\": 100, \"output_dim\": 64},\n",
    "    \"pooling_args\": {\"pooling_method\": \"scatter_sum\"},\n",
    "    \"edge_initialize\": {\"units\": 128, \"use_bias\": True, \"activation\": \"relu\"},\n",
    "    \"edge_dense\": {\"units\": 128, \"use_bias\": True, \"activation\": \"linear\"},\n",
    "    \"edge_activation\": {\"activation\": \"relu\"},\n",
    "    \"node_dense\": {\"units\": 128, \"use_bias\": True, \"activation\": \"relu\"},\n",
    "    \"verbose\": 10, \"depth\": 5,\n",
    "    \"dropout\": {\"rate\": 0.1},\n",
    "    \"output_embedding\": \"graph\",\n",
    "    \"output_mlp\": {\n",
    "        \"use_bias\": [True, True, False], \"units\": [64, 32, 1],\n",
    "        \"activation\": [\"relu\", \"relu\", \"linear\"]\n",
    "    }\n",
    "}\n",
    "# Test making model\n",
    "model = make_model(**model_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b13c8282",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:kgcnn.data.ESOL:Property 'edge_attributes' is an empty list for graph '934'.\n",
      "INFO:kgcnn.data.ESOL:Property 'edge_indices' is an empty list for graph '934'.\n",
      "INFO:kgcnn.data.ESOL:Property 'edge_indices_reverse' is an empty list for graph '934'.\n",
      "WARNING:kgcnn.data.ESOL:Found invalid graphs for properties. Removing graphs '[934]'.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([934])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.clean(model_config[\"inputs\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "74ba8493",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1127, 1)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = np.expand_dims(dataset.get(\"graph_labels\"), axis=-1)\n",
    "labels.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cefd66a",
   "metadata": {},
   "source": [
    "## 3. Keras Tuner\n",
    "\n",
    "Optimize hyperparameter with keras tuner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "435dcaa8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend\n"
     ]
    },
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'tensorflow.keras'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[7], line 2\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mwarnings\u001b[39;00m\n\u001b[1;32m----> 2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mkt\u001b[39;00m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mcopy\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m deepcopy\n\u001b[0;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01msklearn\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmodel_selection\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m train_test_split\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\gcnn_keras_test\\lib\\site-packages\\keras_tuner\\__init__.py:8\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;124;03m\"\"\"DO NOT EDIT.\u001b[39;00m\n\u001b[0;32m      2\u001b[0m \n\u001b[0;32m      3\u001b[0m \u001b[38;5;124;03mThis file was autogenerated. Do not edit it by hand,\u001b[39;00m\n\u001b[0;32m      4\u001b[0m \u001b[38;5;124;03msince your modifications would be overwritten.\u001b[39;00m\n\u001b[0;32m      5\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m----> 8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m applications\n\u001b[0;32m      9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m engine\n\u001b[0;32m     10\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m errors\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\gcnn_keras_test\\lib\\site-packages\\keras_tuner\\applications\\__init__.py:8\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;124;03m\"\"\"DO NOT EDIT.\u001b[39;00m\n\u001b[0;32m      2\u001b[0m \n\u001b[0;32m      3\u001b[0m \u001b[38;5;124;03mThis file was autogenerated. Do not edit it by hand,\u001b[39;00m\n\u001b[0;32m      4\u001b[0m \u001b[38;5;124;03msince your modifications would be overwritten.\u001b[39;00m\n\u001b[0;32m      5\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m----> 8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapplications\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01maugment\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m HyperImageAugment\n\u001b[0;32m      9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapplications\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mefficientnet\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m HyperEfficientNet\n\u001b[0;32m     10\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapplications\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mresnet\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m HyperResNet\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\gcnn_keras_test\\lib\\site-packages\\keras_tuner\\src\\__init__.py:16\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;66;03m# Copyright 2019 The KerasTuner Authors\u001b[39;00m\n\u001b[0;32m      2\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;66;03m# Licensed under the Apache License, Version 2.0 (the \"License\");\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     12\u001b[0m \u001b[38;5;66;03m# See the License for the specific language governing permissions and\u001b[39;00m\n\u001b[0;32m     13\u001b[0m \u001b[38;5;66;03m# limitations under the License.\u001b[39;00m\n\u001b[1;32m---> 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m applications\n\u001b[0;32m     17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m oracles\n\u001b[0;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m tuners\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\gcnn_keras_test\\lib\\site-packages\\keras_tuner\\src\\applications\\__init__.py:16\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;66;03m# Copyright 2019 The KerasTuner Authors\u001b[39;00m\n\u001b[0;32m      2\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;66;03m# Licensed under the Apache License, Version 2.0 (the \"License\");\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     12\u001b[0m \u001b[38;5;66;03m# See the License for the specific language governing permissions and\u001b[39;00m\n\u001b[0;32m     13\u001b[0m \u001b[38;5;66;03m# limitations under the License.\u001b[39;00m\n\u001b[1;32m---> 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapplications\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01maugment\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m HyperImageAugment\n\u001b[0;32m     17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapplications\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mefficientnet\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m HyperEfficientNet\n\u001b[0;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapplications\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mresnet\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m HyperResNet\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\gcnn_keras_test\\lib\\site-packages\\keras_tuner\\src\\applications\\augment.py:17\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;66;03m# Copyright 2019 The KerasTuner Authors\u001b[39;00m\n\u001b[0;32m      2\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;66;03m# Licensed under the Apache License, Version 2.0 (the \"License\");\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     12\u001b[0m \u001b[38;5;66;03m# See the License for the specific language governing permissions and\u001b[39;00m\n\u001b[0;32m     13\u001b[0m \u001b[38;5;66;03m# limitations under the License.\u001b[39;00m\n\u001b[0;32m     16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapi_export\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m keras_tuner_export\n\u001b[1;32m---> 17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackend\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m keras\n\u001b[0;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackend\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ops\n\u001b[0;32m     19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackend\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m random\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\gcnn_keras_test\\lib\\site-packages\\keras_tuner\\src\\backend\\__init__.py:29\u001b[0m\n\u001b[0;32m     27\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackend\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m config\n\u001b[0;32m     28\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackend\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m io\n\u001b[1;32m---> 29\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackend\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m keras\n\u001b[0;32m     30\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackend\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ops\n\u001b[0;32m     31\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkeras_tuner\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msrc\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackend\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m random\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\gcnn_keras_test\\lib\\site-packages\\keras_tuner\\src\\backend\\keras.py:23\u001b[0m\n\u001b[0;32m     21\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m     22\u001b[0m     \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mtf\u001b[39;00m\n\u001b[1;32m---> 23\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mkeras\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m  \u001b[38;5;66;03m# noqa: F403, F401\u001b[39;00m\n\u001b[0;32m     25\u001b[0m     \u001b[38;5;66;03m# Shims to handle symbol renames for older `tf.keras` versions.\u001b[39;00m\n\u001b[0;32m     26\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(tf\u001b[38;5;241m.\u001b[39mkeras, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msaving\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'tensorflow.keras'"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "import keras_tuner as kt\n",
    "from copy import deepcopy\n",
    "from sklearn.model_selection import train_test_split\n",
    "from keras.optimizers import Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6ece2b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_index, test_index = train_test_split(np.arange(len(dataset)), test_size=0.25, random_state=42)\n",
    "x_train, y_train = dataset[train_index].tensor(model_config[\"inputs\"]), labels[train_index]\n",
    "x_valid, y_valid = dataset[test_index].tensor(model_config[\"inputs\"]), labels[test_index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23f06d4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define model build()\n",
    "def build_model(hp):\n",
    "    # Clear clutter from previous Keras session graphs.\n",
    "    clear_session()\n",
    "\n",
    "    hyper_trial = deepcopy(model_config)    \n",
    "    hyper_trial[\"depth\"] = hp.Int('depth', min_value=3, max_value=5, step=1)\n",
    "    hyper_trial[\"pooling_args\"][\"pooling_method\"] = hp.Choice(\"pooling_method\", [\"scatter_sum\", \"scatter_mean\"])\n",
    "    units = hp.Int('nn_size', min_value=25, max_value=400, step=25)\n",
    "    hyper_trial[\"node_dense\"][\"units\"] = units\n",
    "    hyper_trial[\"edge_dense\"][\"units\"] = units\n",
    "    hyper_trial[\"edge_initialize\"][\"units\"] = units\n",
    "    # print(hyper_trial)\n",
    "    \n",
    "    model = make_model(**hyper_trial)\n",
    "    \n",
    "    # We compile our model\n",
    "    learning_rate = hp.Choice('lr_start', [1e-3, 5e-4, 1e-4])\n",
    "    model.compile(\n",
    "        loss=\"mean_squared_error\",\n",
    "        optimizer=Adam(learning_rate=learning_rate),\n",
    "        metrics=[\"mean_absolute_error\"],\n",
    "    )\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5396b881",
   "metadata": {},
   "outputs": [],
   "source": [
    "tuner = kt.Hyperband(build_model,\n",
    "                     objective='val_loss',\n",
    "                     max_epochs=10, factor=3, directory=\"kt_dmpnn\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae7ba82e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tuner.search(x_train,y_train,\n",
    "            shuffle=True,\n",
    "            epochs=100, validation_data=(x_valid, y_valid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c35fcfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_hps = tuner.get_best_hyperparameters(num_trials=1)[0]\n",
    "\n",
    "print(f\"\"\"\n",
    "    Search complete, best hyperparameters:\n",
    "    depth = {best_hps.get('depth')}\n",
    "    nn_size = {best_hps.get('nn_size')}\n",
    "    ene_wt = {best_hps.get('pooling_method')}\n",
    "    learn rate = {best_hps.get('lr_start')}\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6278087-9f9e-45e2-815e-8e1360311a55",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}