mpnn-daint-gpu.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MPNN + GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.6.12 (default, Nov 25 2020, 20:33:10) [GCC]\n",
      "/users/rketkaew/py38-mpnn/bin/python\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "print(sys.version)\n",
    "print(sys.executable)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!module load daint-gpu\n",
    "#!module load TensorFlow/2.4.0-CrayGNU-20.11"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import networkx\n",
    "import ase\n",
    "import ase.visualize\n",
    "\n",
    "from sklearn.utils import shuffle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-3-97ecbf874269>:2: is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.config.list_physical_devices('GPU')` instead.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.test.is_gpu_available(\n",
    "    cuda_only=False, min_cuda_compute_capability=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_inp = \"test/j-coupling-dataset/\"\n",
    "nodes_train = np.load(dir_inp + \"nodes_train.npz\")['arr_0'][:5000]\n",
    "in_edges_train = np.load(dir_inp + \"in_edges_train.npz\")['arr_0'][:5000]\n",
    "out_edges_train = np.load(dir_inp + \"out_edges_train.npz\")['arr_0'][:5000]\n",
    "\n",
    "nodes_test = np.load(dir_inp+ \"nodes_test.npz\")['arr_0'][:1000]\n",
    "in_edges_test = np.load(dir_inp + \"in_edges_test.npz\")['arr_0'][:1000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(5000, 29, 5)\n",
      "(5000, 29, 29, 15)\n",
      "(5000, 29, 29, 1)\n",
      "(1000, 29, 5)\n",
      "(1000, 29, 29, 15)\n"
     ]
    }
   ],
   "source": [
    "print(nodes_train.shape)\n",
    "print(in_edges_train.shape)\n",
    "print(out_edges_train.shape)\n",
    "print(nodes_test.shape)\n",
    "print(in_edges_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_labels = out_edges_train.reshape(-1, out_edges_train.shape[1]*out_edges_train.shape[2],1)\n",
    "in_edges_train = in_edges_train.reshape(-1, in_edges_train.shape[1]*in_edges_train.shape[2], in_edges_train.shape[3])\n",
    "in_edges_test  = in_edges_test.reshape(-1, in_edges_test.shape[1]*in_edges_test.shape[2], in_edges_test.shape[3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes_train, in_edges_train, out_labels = shuffle(nodes_train, in_edges_train, out_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Message Passing Neural Network\n",
    "\n",
    "Implement according to Gilmer et al. https://arxiv.org/abs/1704.01212"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build message parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Message_Passer_NNM(tf.keras.layers.Layer):\n",
    "    def __init__(self, node_dim):\n",
    "        super(Message_Passer_NNM, self).__init__()\n",
    "        self.node_dim = node_dim\n",
    "        self.nn = tf.keras.layers.Dense(units=self.node_dim*self.node_dim, activation = tf.nn.relu)\n",
    "      \n",
    "    def call(self, node_j, edge_ij):\n",
    "        # Embed the edge as a matrix\n",
    "        A = self.nn(edge_ij)\n",
    "        \n",
    "        # Reshape so matrix mult can be done\n",
    "        A = tf.reshape(A, [-1, self.node_dim, self.node_dim])\n",
    "        node_j = tf.reshape(node_j, [-1, self.node_dim, 1])\n",
    "        \n",
    "        # Multiply edge matrix by node and shape into message list\n",
    "        messages = tf.linalg.matmul(A, node_j)\n",
    "        messages = tf.reshape(messages, [-1, tf.shape(edge_ij)[1], self.node_dim])\n",
    "\n",
    "        return messages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build aggregator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Message_Agg(tf.keras.layers.Layer):\n",
    "    def __init__(self):\n",
    "        super(Message_Agg, self).__init__()\n",
    "    \n",
    "    def call(self, messages):\n",
    "        return tf.math.reduce_sum(messages, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build update function - GRU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Update_Func_GRU(tf.keras.layers.Layer):\n",
    "    def __init__(self, state_dim):\n",
    "        super(Update_Func_GRU, self).__init__()\n",
    "        self.concat_layer = tf.keras.layers.Concatenate(axis=1)\n",
    "        self.GRU = tf.keras.layers.GRU(state_dim)\n",
    "        \n",
    "    def call(self, old_state, agg_messages):\n",
    "        # Remember node dim\n",
    "        n_nodes  = tf.shape(old_state)[1]\n",
    "        node_dim = tf.shape(old_state)[2]\n",
    "        \n",
    "        # Reshape so GRU can be applied, concat so old_state and messages are in sequence\n",
    "        old_state = tf.reshape(old_state, [-1, 1, tf.shape(old_state)[-1]])\n",
    "        agg_messages = tf.reshape(agg_messages, [-1, 1, tf.shape(agg_messages)[-1]])\n",
    "        concat = self.concat_layer([old_state, agg_messages])\n",
    "        \n",
    "        # Apply GRU and then reshape so it can be returned\n",
    "        activation = self.GRU(concat)\n",
    "        activation = tf.reshape(activation, [-1, n_nodes, node_dim])\n",
    "        \n",
    "        return activation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Output layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Edge_Regressor(tf.keras.layers.Layer):\n",
    "    def __init__(self, intermediate_dim):\n",
    "        super(Edge_Regressor, self).__init__()\n",
    "        self.concat_layer = tf.keras.layers.Concatenate()\n",
    "        self.hidden_layer_1 = tf.keras.layers.Dense(units=intermediate_dim, activation=tf.nn.relu)\n",
    "        self.hidden_layer_2 = tf.keras.layers.Dense(units=intermediate_dim, activation=tf.nn.relu)\n",
    "        self.output_layer = tf.keras.layers.Dense(units=1, activation=None)\n",
    "\n",
    "    def call(self, nodes, edges):\n",
    "        # Remember node dims\n",
    "        n_nodes  = tf.shape(nodes)[1]\n",
    "        node_dim = tf.shape(nodes)[2]\n",
    "        \n",
    "        # Tile and reshape to match edges\n",
    "        state_i = tf.reshape(tf.tile(nodes, [1, 1, n_nodes]),[-1,n_nodes*n_nodes, node_dim ])\n",
    "        state_j = tf.tile(nodes, [1, n_nodes, 1])\n",
    "        \n",
    "        # concat edges and nodes and apply MLP\n",
    "        concat = self.concat_layer([state_i, edges, state_j])\n",
    "        activation_1 = self.hidden_layer_1(concat)  \n",
    "        activation_2 = self.hidden_layer_2(activation_1)\n",
    "\n",
    "        return self.output_layer(activation_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build Single Message Passing Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MP_Layer(tf.keras.layers.Layer):\n",
    "    def __init__(self, state_dim):\n",
    "        super(MP_Layer, self).__init__(self)\n",
    "        self.message_passers  = Message_Passer_NNM(node_dim = state_dim) \n",
    "        self.message_aggs    = Message_Agg()\n",
    "        self.update_functions = Update_Func_GRU(state_dim = state_dim)\n",
    "        self.state_dim = state_dim         \n",
    "\n",
    "    def call(self, nodes, edges, mask):\n",
    "        n_nodes  = tf.shape(nodes)[1]\n",
    "        node_dim = tf.shape(nodes)[2]\n",
    "        \n",
    "        state_j = tf.tile(nodes, [1, n_nodes, 1])\n",
    "\n",
    "        messages  = self.message_passers(state_j, edges)\n",
    "\n",
    "        # Do this to ignore messages from non-existant nodes\n",
    "        masked =  tf.math.multiply(messages, mask)\n",
    "        \n",
    "        masked = tf.reshape(masked, [tf.shape(messages)[0], n_nodes, n_nodes, node_dim])\n",
    "\n",
    "        agg_m = self.message_aggs(masked)\n",
    "        \n",
    "        updated_nodes = self.update_functions(nodes, agg_m)\n",
    "        \n",
    "        nodes_out = updated_nodes\n",
    "        # Batch norm seems not to work. \n",
    "        #nodes_out = self.batch_norm(updated_nodes)\n",
    "        \n",
    "        return nodes_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Formulate MPNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "adj_input = tf.keras.Input(shape=(None,), name='adj_input')\n",
    "nod_input = tf.keras.Input(shape=(None,), name='nod_input')\n",
    "\n",
    "class MPNN(tf.keras.Model):\n",
    "    def __init__(self, out_int_dim, state_dim, T):\n",
    "        super(MPNN, self).__init__(self)   \n",
    "        self.T = T\n",
    "        self.embed = tf.keras.layers.Dense(units=state_dim, activation=tf.nn.relu)\n",
    "        self.MP = MP_Layer( state_dim)     \n",
    "        self.edge_regressor  = Edge_Regressor(out_int_dim)\n",
    "        #self.batch_norm = tf.keras.layers.BatchNormalization()\n",
    "        \n",
    "    def call(self, inputs =  [adj_input, nod_input]):\n",
    "        nodes            = inputs['nod_input']\n",
    "        edges            = inputs['adj_input']\n",
    "\n",
    "        # Get distances, and create mask wherever 0 (i.e. non-existant nodes)\n",
    "        # This also masks node self-interactions...\n",
    "        # This assumes distance is last\n",
    "        len_edges = tf.shape(edges)[-1]\n",
    "        \n",
    "        _, x = tf.split(edges, [len_edges -1, 1], 2)\n",
    "        mask =  tf.where(tf.equal(x, 0), x, tf.ones_like(x))\n",
    "        \n",
    "        # Embed node to be of the chosen node dimension (you can also just pad)\n",
    "        nodes = self.embed(nodes) \n",
    "        \n",
    "        #nodes = self.batch_norm(nodes)\n",
    "        # Run the T message passing steps\n",
    "        for mp in range(self.T):\n",
    "            nodes =  self.MP(nodes, edges, mask)\n",
    "        \n",
    "        # Regress the output values\n",
    "        con_edges = self.edge_regressor(nodes, edges)\n",
    "        \n",
    "        return con_edges"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define metrics (loss)\n",
    "\n",
    "Supported now:\n",
    "- MSE\n",
    "- Log MSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mse(orig , preds):\n",
    "    # Mask values for which no scalar coupling exists\n",
    "    mask  = tf.where(tf.equal(orig, 0), orig, tf.ones_like(orig))\n",
    "\n",
    "    nums  = tf.boolean_mask(orig,  mask)\n",
    "    preds = tf.boolean_mask(preds,  mask)\n",
    "\n",
    "    reconstruction_error = tf.reduce_mean(tf.square(tf.subtract(nums, preds)))\n",
    "\n",
    "    return reconstruction_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_mse(orig , preds):\n",
    "    # Mask values for which no scalar coupling exists\n",
    "    mask  = tf.where(tf.equal(orig, 0), orig, tf.ones_like(orig))\n",
    "\n",
    "    nums  = tf.boolean_mask(orig,  mask)\n",
    "    preds = tf.boolean_mask(preds,  mask)\n",
    "\n",
    "    reconstruction_error = tf.math.log(tf.reduce_mean(tf.square(tf.subtract(nums, preds))))\n",
    "\n",
    "    return reconstruction_error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define callback and optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 0.001\n",
    "def step_decay(epoch):\n",
    "    initial_lrate = learning_rate\n",
    "    drop = 0.1\n",
    "    epochs_drop = 20.0\n",
    "    lrate = initial_lrate * np.power(drop,  \n",
    "           np.floor((epoch)/epochs_drop))\n",
    "    tf.print(\"Learning rate: \", lrate)\n",
    "    return lrate\n",
    "\n",
    "lrate = tf.keras.callbacks.LearningRateScheduler(step_decay)\n",
    "stop_early = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience = 15, restore_best_weights=True)\n",
    "\n",
    "#lrate  =  tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1,\n",
    "#                              patience=5, min_lr=0.00001, verbose = 1)\n",
    "\n",
    "opt = tf.optimizers.Adam(learning_rate=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct a model and compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10, 841, 1), dtype=float32, numpy=\n",
       "array([[[ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]],\n",
       "\n",
       "       [[ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]],\n",
       "\n",
       "       [[ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]],\n",
       "\n",
       "       [[-0.00705141],\n",
       "        [-0.1289805 ],\n",
       "        [-0.02763902],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]],\n",
       "\n",
       "       [[ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        ...,\n",
       "        [ 0.        ],\n",
       "        [ 0.        ],\n",
       "        [ 0.        ]]], dtype=float32)>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mpnn = MPNN(out_int_dim = 512, state_dim = 128, T = 4)\n",
    "mpnn.compile(opt, mse, metrics = [mse, log_mse])\n",
    "\n",
    "train_size = int(len(out_labels)*0.8)\n",
    "batch_size = 16\n",
    "epochs = 25\n",
    "\n",
    "mpnn.call({'adj_input' : in_edges_train[:10], 'nod_input': nodes_train[:10]})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/25\n",
      "Learning rate:  0.001\n",
      "143/250 [================>.............] - ETA: 2:14 - loss: nan - mse: nan - log_mse: nan"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-29-57018283b957>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     10\u001b[0m          validation_data = ({'adj_input' : in_edges_train[train_size:], \n\u001b[1;32m     11\u001b[0m                              'nod_input': nodes_train[train_size:]}, \n\u001b[0;32m---> 12\u001b[0;31m                             out_labels[train_size:]))\n\u001b[0m",
      "\u001b[0;32m~/py38-mpnn/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)\u001b[0m\n\u001b[1;32m   1181\u001b[0m                 _r=1):\n\u001b[1;32m   1182\u001b[0m               \u001b[0mcallbacks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_train_batch_begin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1183\u001b[0;31m               \u001b[0mtmp_logs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0miterator\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1184\u001b[0m               \u001b[0;32mif\u001b[0m \u001b[0mdata_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshould_sync\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1185\u001b[0m                 \u001b[0mcontext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0masync_wait\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/py38-mpnn/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    887\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    888\u001b[0m       \u001b[0;32mwith\u001b[0m \u001b[0mOptionalXlaContext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_jit_compile\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 889\u001b[0;31m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    890\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    891\u001b[0m       \u001b[0mnew_tracing_count\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/py38-mpnn/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36m_call\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    915\u001b[0m       \u001b[0;31m# In this case we have created variables on the first call, so we run the\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    916\u001b[0m       \u001b[0;31m# defunned version which is guaranteed to never create variables.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 917\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stateless_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# pylint: disable=not-callable\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    918\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stateful_fn\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    919\u001b[0m       \u001b[0;31m# Release the lock early so that multiple threads can perform the call\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/py38-mpnn/lib/python3.6/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   3022\u001b[0m        filtered_flat_args) = self._maybe_define_function(args, kwargs)\n\u001b[1;32m   3023\u001b[0m     return graph_function._call_flat(\n\u001b[0;32m-> 3024\u001b[0;31m         filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access\n\u001b[0m\u001b[1;32m   3025\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3026\u001b[0m   \u001b[0;34m@\u001b[0m\u001b[0mproperty\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/py38-mpnn/lib/python3.6/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_call_flat\u001b[0;34m(self, args, captured_inputs, cancellation_manager)\u001b[0m\n\u001b[1;32m   1959\u001b[0m       \u001b[0;31m# No tape is watching; skip to running the function.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1960\u001b[0m       return self._build_call_outputs(self._inference_function.call(\n\u001b[0;32m-> 1961\u001b[0;31m           ctx, args, cancellation_manager=cancellation_manager))\n\u001b[0m\u001b[1;32m   1962\u001b[0m     forward_backward = self._select_forward_and_backward_functions(\n\u001b[1;32m   1963\u001b[0m         \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/py38-mpnn/lib/python3.6/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36mcall\u001b[0;34m(self, ctx, args, cancellation_manager)\u001b[0m\n\u001b[1;32m    594\u001b[0m               \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    595\u001b[0m               \u001b[0mattrs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mattrs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 596\u001b[0;31m               ctx=ctx)\n\u001b[0m\u001b[1;32m    597\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    598\u001b[0m           outputs = execute.execute_with_cancellation(\n",
      "\u001b[0;32m~/py38-mpnn/lib/python3.6/site-packages/tensorflow/python/eager/execute.py\u001b[0m in \u001b[0;36mquick_execute\u001b[0;34m(op_name, num_outputs, inputs, attrs, ctx, name)\u001b[0m\n\u001b[1;32m     58\u001b[0m     \u001b[0mctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mensure_initialized\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     59\u001b[0m     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,\n\u001b[0;32m---> 60\u001b[0;31m                                         inputs, attrs, num_outputs)\n\u001b[0m\u001b[1;32m     61\u001b[0m   \u001b[0;32mexcept\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_NotOkStatusException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     62\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "mpnn.fit({'adj_input': in_edges_train[:train_size], \n",
    "          'nod_input': nodes_train[:train_size]}, \n",
    "         y = out_labels[:train_size], \n",
    "         batch_size = batch_size, \n",
    "         epochs = epochs, \n",
    "         callbacks = [lrate, stop_early], \n",
    "         use_multiprocessing = True, \n",
    "         initial_epoch = 0, \n",
    "         verbose = 1, \n",
    "         validation_data = ({'adj_input' : in_edges_train[train_size:], \n",
    "                             'nod_input': nodes_train[train_size:]}, \n",
    "                            out_labels[train_size:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Prediction\n",
    "\n",
    "preds = mpnn.predict({'adj_input' : in_edges_test, 'nod_input': nodes_test})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('py38')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "f2b0cf5112dfd2af73b780f4815bf58f6493434a6fee69d3bee328e26a09a438"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}