FinalProject.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "RDKit WARNING: [23:42:54] Enabling RDKit 2019.09.1 jupyter extensions\n"
     ]
    }
   ],
   "source": [
    "from __future__ import print_function\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "import time\n",
    "import argparse\n",
    "import os,sys\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import random\n",
    "import argparse\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import ChemicalFeatures\n",
    "from rdkit import RDConfig\n",
    "from joblib import Parallel, delayed\n",
    "from os import listdir\n",
    "from os.path import isfile, join\n",
    "import xml.etree.ElementTree as ET\n",
    "import shutil\n",
    "import torch.utils.data as data\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd.variable import Variable\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import ChemicalFeatures\n",
    "from rdkit import RDConfig\n",
    "from os import listdir\n",
    "import rdkit\n",
    "import multiprocessing\n",
    "from os.path import isfile, join\n",
    "from datasets import utils\n",
    "random.seed(2)\n",
    "np.random.seed(2)\n",
    "best_er1=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description='Neural message passing')\n",
    "parser.add_argument('--root', nargs=1, help='Specify the data directory.', default=['data-parse/chemfindata/'])\n",
    "parser.add_argument('--dataset', default='our_data')\n",
    "parser.add_argument('--datasetPath', default='./data-parse/chemfindata/', help='dataset path')\n",
    "parser.add_argument('--logPath', default='./log/our_data/mpnn/', help='log path')\n",
    "parser.add_argument('--plotLr', default=False, help='allow plotting the data')\n",
    "parser.add_argument('--plotPath', default='./plot/our_data/mpnn/', help='plot path')\n",
    "parser.add_argument('--resume', default='./checkpoint/our_data/mpnn/',\n",
    "                    help='path to latest checkpoint')\n",
    "parser.add_argument('--batch-size', type=int, default=10, metavar='N',\n",
    "                    help='Input batch size for training (default: 20)')\n",
    "parser.add_argument('--no-cuda', action='store_true', default=False,\n",
    "                    help='Enables CUDA training')\n",
    "parser.add_argument('--epochs', type=int, default=360, metavar='N',\n",
    "                    help='Number of epochs to train (default: 360)')\n",
    "parser.add_argument('--lr', type=lambda x: restricted_float(x, [1e-5, 1e-2]), default=1e-3, metavar='LR',\n",
    "                    help='Initial learning rate [1e-5, 5e-4] (default: 1e-4)')\n",
    "parser.add_argument('--lr-decay', type=lambda x: restricted_float(x, [.01, 1]), default=0.6, metavar='LR-DECAY',\n",
    "                    help='Learning rate decay factor [.01, 1] (default: 0.6)')\n",
    "parser.add_argument('--schedule', type=list, default=[0.1, 0.9], metavar='S',\n",
    "                    help='Percentage of epochs to start the learning rate decay [0, 1] (default: [0.1, 0.9])')\n",
    "parser.add_argument('--momentum', type=float, default=0.9, metavar='M',\n",
    "                    help='SGD momentum (default: 0.9)')\n",
    "parser.add_argument('--log-interval', type=int, default=20, metavar='N',\n",
    "                    help='How many batches to wait before logging training status')\n",
    "parser.add_argument('--prefetch', type=int, default=2, help='Pre-fetching threads.')\n",
    "args = parser.parse_args([])\n",
    "args.cuda = not args.no_cuda and torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NNet(nn.Module):\n",
    "\n",
    "    def __init__(self, n_in, n_out, hlayers=(128, 256, 128)):\n",
    "        super(NNet, self).__init__()\n",
    "        self.n_hlayers = len(hlayers)\n",
    "        self.fcs = nn.ModuleList([nn.Linear(n_in, hlayers[i]) if i == 0 else\n",
    "                                  nn.Linear(hlayers[i-1], n_out) if i == self.n_hlayers else\n",
    "                                  nn.Linear(hlayers[i-1], hlayers[i]) for i in range(self.n_hlayers+1)])\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.contiguous().view(-1, self.num_flat_features(x))\n",
    "        for i in range(self.n_hlayers):\n",
    "            x = F.relu(self.fcs[i](x))\n",
    "        x = self.fcs[-1](x)\n",
    "        return x\n",
    "\n",
    "    def num_flat_features(self, x):\n",
    "        size = x.size()[1:]  # all dimensions except the batch dimension\n",
    "        num_features = 1\n",
    "        for s in size:\n",
    "            num_features *= s\n",
    "        return num_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MessageFunction(nn.Module):\n",
    "\n",
    "    # Constructor\n",
    "    def __init__(self, message_def='duvenaud', args={}):\n",
    "        super(MessageFunction, self).__init__()\n",
    "        self.m_definition = ''\n",
    "        self.m_function = None\n",
    "        self.args = {}\n",
    "        self.__set_message(message_def, args)\n",
    "\n",
    "    # Message from h_v to h_w through e_vw\n",
    "    def forward(self, h_v, h_w, e_vw, args=None):\n",
    "        return self.m_function(h_v, h_w, e_vw, args)\n",
    "\n",
    "    # Set a message function\n",
    "    def __set_message(self, message_def, args={}):\n",
    "        self.m_definition = message_def.lower()\n",
    "\n",
    "        self.m_function = {\n",
    "                    'duvenaud':         self.m_duvenaud,\n",
    "                    'ggnn':             self.m_ggnn,\n",
    "                    'intnet':           self.m_intnet,\n",
    "                    'mpnn':             self.m_mpnn,\n",
    "                    'mgc':              self.m_mgc,\n",
    "                    'bruna':            self.m_bruna,\n",
    "                    'defferrard':       self.m_deff,\n",
    "                    'kipf':             self.m_kipf\n",
    "                }.get(self.m_definition, None)\n",
    "\n",
    "        if self.m_function is None:\n",
    "            print('WARNING!: Message Function has not been set correctly\\n\\tIncorrect definition ' + message_def)\n",
    "            quit()\n",
    "\n",
    "        init_parameters = {\n",
    "            'duvenaud': self.init_duvenaud,\n",
    "            'ggnn':     self.init_ggnn,\n",
    "            'intnet':   self.init_intnet,\n",
    "            'mpnn':     self.init_mpnn\n",
    "        }.get(self.m_definition, lambda x: (nn.ParameterList([]), nn.ModuleList([]), {}))\n",
    "\n",
    "        self.learn_args, self.learn_modules, self.args = init_parameters(args)\n",
    "\n",
    "        self.m_size = {\n",
    "                'duvenaud':     self.out_duvenaud,\n",
    "                'ggnn':         self.out_ggnn,\n",
    "                'intnet':       self.out_intnet,\n",
    "                'mpnn':         self.out_mpnn\n",
    "            }.get(self.m_definition, None)\n",
    "\n",
    "    # Get the name of the used message function\n",
    "    def get_definition(self):\n",
    "        return self.m_definition\n",
    "\n",
    "    # Get the message function arguments\n",
    "    def get_args(self):\n",
    "        return self.args\n",
    "\n",
    "    # Get Output size\n",
    "    def get_out_size(self, size_h, size_e, args=None):\n",
    "        return self.m_size(size_h, size_e, args)\n",
    "\n",
    "    # Definition of various state of the art message functions\n",
    "    \n",
    "    # Duvenaud et al. (2015), Convolutional Networks for Learning Molecular Fingerprints\n",
    "    def m_duvenaud(self, h_v, h_w, e_vw, args):\n",
    "        m = torch.cat([h_w, e_vw], 2)\n",
    "        return m\n",
    "\n",
    "    def out_duvenaud(self, size_h, size_e, args):\n",
    "        return size_h + size_e\n",
    "\n",
    "    def init_duvenaud(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n",
    "\n",
    "    # Li et al. (2016), Gated Graph Neural Networks (GG-NN)\n",
    "    def m_ggnn(self, h_v, h_w, e_vw, opt={}):\n",
    "\n",
    "        m = Variable(torch.zeros(h_w.size(0), h_w.size(1), self.args['out']).type_as(h_w.data))\n",
    "\n",
    "        for w in range(h_w.size(1)):\n",
    "            if torch.nonzero(e_vw[:, w, :].data).size():\n",
    "                for i, el in enumerate(self.args['e_label']):\n",
    "                    ind = (el == e_vw[:,w,:]).type_as(self.learn_args[0][i])\n",
    "\n",
    "                    parameter_mat = self.learn_args[0][i][None, ...].expand(h_w.size(0), self.learn_args[0][i].size(0),\n",
    "                                                                            self.learn_args[0][i].size(1))\n",
    "\n",
    "                    m_w = torch.transpose(torch.bmm(torch.transpose(parameter_mat, 1, 2),\n",
    "                                                                        torch.transpose(torch.unsqueeze(h_w[:, w, :], 1),\n",
    "                                                                                        1, 2)), 1, 2)\n",
    "                    m_w = torch.squeeze(m_w)\n",
    "                    m[:,w,:] = ind.expand_as(m_w)*m_w\n",
    "        return m\n",
    "\n",
    "    def out_ggnn(self, size_h, size_e, args):\n",
    "        return self.args['out']\n",
    "\n",
    "    def init_ggnn(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "\n",
    "        args['e_label'] = params['e_label']\n",
    "        args['in'] = params['in']\n",
    "        args['out'] = params['out']\n",
    "\n",
    "        # Define a parameter matrix A for each edge label.\n",
    "        learn_args.append(nn.Parameter(torch.randn(len(params['e_label']), params['in'], params['out'])))\n",
    "\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n",
    "\n",
    "    # Battaglia et al. (2016), Interaction Networks\n",
    "    def m_intnet(self, h_v, h_w, e_vw, args):\n",
    "        m = torch.cat([h_v[:, None, :].expand_as(h_w), h_w, e_vw], 2)\n",
    "        b_size = m.size()\n",
    "\n",
    "        m = m.view(-1, b_size[2])\n",
    "\n",
    "        m = self.learn_modules[0](m)\n",
    "        m = m.view(b_size[0], b_size[1], -1)\n",
    "        return m\n",
    "\n",
    "    def out_intnet(self, size_h, size_e, args):\n",
    "        return self.args['out']\n",
    "\n",
    "    def init_intnet(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "        args['in'] = params['in']\n",
    "        args['out'] = params['out']\n",
    "        learn_modules.append(NNet(n_in=params['in'], n_out=params['out']))\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n",
    "\n",
    "    # Gilmer et al. (2017), Neural Message Passing for Quantum Chemistry\n",
    "    def m_mpnn(self, h_v, h_w, e_vw, opt={}):\n",
    "        # Matrices for each edge\n",
    "        edge_output = self.learn_modules[0](e_vw)\n",
    "        edge_output = edge_output.view(-1, self.args['out'], self.args['in'])\n",
    "\n",
    "        h_w_rows = h_w[..., None].expand(h_w.size(0), h_w.size(1), h_v.size(1)).contiguous()\n",
    "\n",
    "        h_w_rows = h_w_rows.view(-1, self.args['in'])\n",
    "\n",
    "        h_multiply = torch.bmm(edge_output, torch.unsqueeze(h_w_rows,2))\n",
    "\n",
    "        m_new = torch.squeeze(h_multiply)\n",
    "\n",
    "        return m_new\n",
    "\n",
    "    def out_mpnn(self, size_h, size_e, args):\n",
    "        return self.args['out']\n",
    "\n",
    "    def init_mpnn(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "\n",
    "        args['in'] = params['in']\n",
    "        args['out'] = params['out']\n",
    "\n",
    "        # Define a parameter matrix A for each edge label.\n",
    "        learn_modules.append(NNet(n_in=params['edge_feat'], n_out=(params['in']*params['out'])))\n",
    "\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n",
    "\n",
    "    # Kearnes et al. (2016), Molecular Graph Convolutions\n",
    "    def m_mgc(self, h_v, h_w, e_vw, args):\n",
    "        m = e_vw\n",
    "        return m\n",
    "    \n",
    "    # Laplacian based methods\n",
    "    # Bruna et al. (2013)\n",
    "    def m_bruna(self, h_v, h_w, e_vw, args):\n",
    "        # TODO\n",
    "        m = [] \n",
    "        return m\n",
    "\n",
    "    # Defferrard et al. (2016)\n",
    "    def m_deff(self, h_v, h_w, e_vw, args):\n",
    "        # TODO\n",
    "        m = []\n",
    "        return m\n",
    "\n",
    "    # Kipf & Welling (2016)\n",
    "    def m_kipf(self, h_v, h_w, e_vw, args):\n",
    "        # TODO\n",
    "        m = []\n",
    "        return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UpdateFunction(nn.Module):\n",
    "\n",
    "    # Constructor\n",
    "    def __init__(self, update_def='nn', args={}):\n",
    "        super(UpdateFunction, self).__init__()\n",
    "        self.u_definition = ''\n",
    "        self.u_function = None\n",
    "        self.args = {}\n",
    "        self.__set_update(update_def, args)\n",
    "\n",
    "    # Update node hv given message mv\n",
    "    def forward(self, h_v, m_v, opt={}):\n",
    "        return self.u_function(h_v, m_v, opt)\n",
    "\n",
    "    # Set update function\n",
    "    def __set_update(self, update_def, args):\n",
    "        self.u_definition = update_def.lower()\n",
    "\n",
    "        self.u_function = {\n",
    "                    'duvenaud':         self.u_duvenaud,\n",
    "                    'ggnn':             self.u_ggnn,\n",
    "                    'intnet':           self.u_intnet,\n",
    "                    'mpnn':             self.u_mpnn\n",
    "                }.get(self.u_definition, None)\n",
    "\n",
    "        if self.u_function is None:\n",
    "            print('WARNING!: Update Function has not been set correctly\\n\\tIncorrect definition ' + update_def)\n",
    "\n",
    "        init_parameters = {\n",
    "            'duvenaud':         self.init_duvenaud,\n",
    "            'ggnn':             self.init_ggnn,\n",
    "            'intnet':           self.init_intnet,\n",
    "            'mpnn':             self.init_mpnn\n",
    "        }.get(self.u_definition, lambda x: (nn.ParameterList([]), nn.ModuleList([]), {}))\n",
    "\n",
    "        self.learn_args, self.learn_modules, self.args = init_parameters(args)\n",
    "\n",
    "    # Get the name of the used update function\n",
    "    def get_definition(self):\n",
    "        return self.u_definition\n",
    "\n",
    "    # Get the update function arguments\n",
    "    def get_args(self):\n",
    "        return self.args\n",
    "\n",
    "    ## Definition of various state of the art update functions\n",
    "\n",
    "    # Duvenaud\n",
    "    def u_duvenaud(self, h_v, m_v, opt):\n",
    "\n",
    "        param_sz = self.learn_args[0][opt['deg']].size()\n",
    "        parameter_mat = torch.t(self.learn_args[0][opt['deg']])[None, ...].expand(m_v.size(0), param_sz[1], param_sz[0])\n",
    "\n",
    "        aux = torch.bmm(parameter_mat, torch.transpose(m_v, 1, 2))\n",
    "\n",
    "        return torch.transpose(torch.nn.Sigmoid()(aux), 1, 2)\n",
    "\n",
    "    def init_duvenaud(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "\n",
    "        # Filter degree 0 (the message will be 0 and therefore there is no update\n",
    "        args['deg'] = [i for i in params['deg'] if i!=0]\n",
    "        args['in'] = params['in']\n",
    "        args['out'] = params['out']\n",
    "\n",
    "        # Define a parameter matrix H for each degree.\n",
    "        learn_args.append(torch.nn.Parameter(torch.randn(len(args['deg']), args['in'], args['out'])))\n",
    "\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n",
    "\n",
    "    # GG-NN, Li et al.\n",
    "    def u_ggnn(self, h_v, m_v, opt={}):\n",
    "        h_v.contiguous()\n",
    "        m_v.contiguous()\n",
    "        h_new = self.learn_modules[0](torch.transpose(m_v, 0, 1), torch.unsqueeze(h_v, 0))[0]  # 0 or 1???\n",
    "        return torch.transpose(h_new, 0, 1)\n",
    "\n",
    "    def init_ggnn(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "\n",
    "        args['in_m'] = params['in_m']\n",
    "        args['out'] = params['out']\n",
    "\n",
    "        # GRU\n",
    "        learn_modules.append(nn.GRU(params['in_m'], params['out']))\n",
    "\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n",
    "\n",
    "    # Battaglia et al. (2016), Interaction Networks\n",
    "    def u_intnet(self, h_v, m_v, opt):\n",
    "        if opt['x_v'].ndimension():\n",
    "            input_tensor = torch.cat([h_v, opt['x_v'], torch.squeeze(m_v)], 1)\n",
    "        else:\n",
    "            input_tensor = torch.cat([h_v, torch.squeeze(m_v)], 1)\n",
    "\n",
    "        return self.learn_modules[0](input_tensor)\n",
    "\n",
    "    def init_intnet(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "\n",
    "        args['in'] = params['in']\n",
    "        args['out'] = params['out']\n",
    "\n",
    "        learn_modules.append(NNet(n_in=params['in'], n_out=params['out']))\n",
    "\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n",
    "\n",
    "    def u_mpnn(self, h_v, m_v, opt={}):\n",
    "        h_in = h_v.view(-1,h_v.size(2))\n",
    "        m_in = m_v.view(-1,m_v.size(2))\n",
    "        h_new = self.learn_modules[0](m_in[None,...],h_in[None,...])[0] # 0 or 1???\n",
    "        return torch.squeeze(h_new).view(h_v.size())\n",
    "\n",
    "    def init_mpnn(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "\n",
    "        args['in_m'] = params['in_m']\n",
    "        args['out'] = params['out']\n",
    "\n",
    "        # GRU\n",
    "        learn_modules.append(nn.GRU(params['in_m'], params['out']))\n",
    "\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReadoutFunction(nn.Module):\n",
    "\n",
    "    # Constructor\n",
    "    def __init__(self, readout_def='nn', args={}):\n",
    "        super(ReadoutFunction, self).__init__()\n",
    "        self.r_definition = ''\n",
    "        self.r_function = None\n",
    "        self.args = {}\n",
    "        self.__set_readout(readout_def, args)\n",
    "\n",
    "    # Readout graph given node values at las layer\n",
    "    def forward(self, h_v):\n",
    "        return self.r_function(h_v)\n",
    "\n",
    "    # Set a readout function\n",
    "    def __set_readout(self, readout_def, args):\n",
    "        self.r_definition = readout_def.lower()\n",
    "\n",
    "        self.r_function = {\n",
    "                    'duvenaud': self.r_duvenaud,\n",
    "                    'ggnn':     self.r_ggnn,\n",
    "                    'intnet':   self.r_intnet,\n",
    "                    'mpnn':     self.r_mpnn\n",
    "                }.get(self.r_definition, None)\n",
    "\n",
    "        if self.r_function is None:\n",
    "            print('WARNING!: Readout Function has not been set correctly\\n\\tIncorrect definition ' + readout_def)\n",
    "            quit()\n",
    "\n",
    "        init_parameters = {\n",
    "            'duvenaud': self.init_duvenaud,\n",
    "            'ggnn':     self.init_ggnn,\n",
    "            'intnet':   self.init_intnet,\n",
    "            'mpnn':     self.init_mpnn\n",
    "        }.get(self.r_definition, lambda x: (nn.ParameterList([]), nn.ModuleList([]), {}))\n",
    "\n",
    "        self.learn_args, self.learn_modules, self.args = init_parameters(args)\n",
    "\n",
    "    # Get the name of the used readout function\n",
    "    def get_definition(self):\n",
    "        return self.r_definition\n",
    "\n",
    "    ## Definition of various state of the art update functions\n",
    "\n",
    "    # Duvenaud\n",
    "    def r_duvenaud(self, h):\n",
    "        # layers\n",
    "        aux = []\n",
    "        for l in range(len(h)):\n",
    "            param_sz = self.learn_args[l].size()\n",
    "            parameter_mat = torch.t(self.learn_args[l])[None, ...].expand(h[l].size(0), param_sz[1],\n",
    "                                                                                      param_sz[0])\n",
    "\n",
    "            aux.append(torch.transpose(torch.bmm(parameter_mat, torch.transpose(h[l], 1, 2)), 1, 2))\n",
    "\n",
    "            for j in range(0, aux[l].size(1)):\n",
    "                # Mask whole 0 vectors\n",
    "                aux[l][:, j, :] = nn.Softmax()(aux[l][:, j, :].clone())*(torch.sum(aux[l][:, j, :] != 0, 1) > 0).expand_as(aux[l][:, j, :]).type_as(aux[l])\n",
    "\n",
    "        aux = torch.sum(torch.sum(torch.stack(aux, 3), 3), 1)\n",
    "        return self.learn_modules[0](torch.squeeze(aux))\n",
    "\n",
    "    def init_duvenaud(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "\n",
    "        args['out'] = params['out']\n",
    "\n",
    "        # Define a parameter matrix W for each layer.\n",
    "        for l in range(params['layers']):\n",
    "            learn_args.append(nn.Parameter(torch.randn(params['in'][l], params['out'])))\n",
    "\n",
    "        # learn_modules.append(nn.Linear(params['out'], params['target']))\n",
    "\n",
    "        learn_modules.append(NNet(n_in=params['out'], n_out=params['target']))\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n",
    "\n",
    "    # GG-NN, Li et al.\n",
    "    def r_ggnn(self, h):\n",
    "\n",
    "        aux = Variable( torch.Tensor(h[0].size(0), self.args['out']).type_as(h[0].data).zero_() )\n",
    "        # For each graph\n",
    "        for i in range(h[0].size(0)):\n",
    "            nn_res = nn.Sigmoid()(self.learn_modules[0](torch.cat([h[0][i,:,:], h[-1][i,:,:]], 1)))*self.learn_modules[1](h[-1][i,:,:])\n",
    "\n",
    "            # Delete virtual nodes\n",
    "            nn_res = (torch.sum(h[0][i,:,:],1).expand_as(nn_res)>0).type_as(nn_res)* nn_res\n",
    "\n",
    "            aux[i,:] = torch.sum(nn_res,0)\n",
    "\n",
    "        return aux\n",
    "\n",
    "    def init_ggnn(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "\n",
    "        # i\n",
    "        learn_modules.append(NNet(n_in=2*params['in'], n_out=params['target']))\n",
    "\n",
    "        # j\n",
    "        learn_modules.append(NNet(n_in=params['in'], n_out=params['target']))\n",
    "\n",
    "        args['out'] = params['target']\n",
    "\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n",
    "\n",
    "\n",
    "    # Battaglia et al. (2016), Interaction Networks\n",
    "    def r_intnet(self, h):\n",
    "\n",
    "        aux = torch.sum(h[-1],1)\n",
    "\n",
    "        return self.learn_modules[0](aux)\n",
    "\n",
    "    def init_intnet(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "\n",
    "        learn_modules.append(NNet(n_in=params['in'], n_out=params['target']))\n",
    "\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n",
    "\n",
    "    def r_mpnn(self, h):\n",
    "\n",
    "        aux = Variable( torch.Tensor(h[0].size(0), self.args['out']).type_as(h[0].data).zero_() )\n",
    "        # For each graph\n",
    "        for i in range(h[0].size(0)):\n",
    "            nn_res = nn.Sigmoid()(self.learn_modules[0](torch.cat([h[0][i,:,:], h[-1][i,:,:]], 1)))*self.learn_modules[1](h[-1][i,:,:])\n",
    "\n",
    "            # Delete virtual nodes\n",
    "            nn_res = (torch.sum(h[0][i,:,:],1)[...,None].expand_as(nn_res)>0).type_as(nn_res)* nn_res\n",
    "\n",
    "            aux[i,:] = torch.sum(nn_res,0)\n",
    "\n",
    "        return aux\n",
    "\n",
    "    def init_mpnn(self, params):\n",
    "        learn_args = []\n",
    "        learn_modules = []\n",
    "        args = {}\n",
    "\n",
    "        # i\n",
    "        learn_modules.append(NNet(n_in=2*params['in'], n_out=params['target']))\n",
    "\n",
    "        # j\n",
    "        learn_modules.append(NNet(n_in=params['in'], n_out=params['target']))\n",
    "\n",
    "        args['out'] = params['target']\n",
    "\n",
    "        return nn.ParameterList(learn_args), nn.ModuleList(learn_modules), args\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MpnnGGNN(nn.Module):\n",
    "    \"\"\"\n",
    "        MPNN as proposed by Li et al..\n",
    "\n",
    "        This class implements the whole Li et al. model following the functions proposed by Gilmer et al. as\n",
    "        Message, Update and Readout.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        e : int list.\n",
    "            Possible edge labels for the input graph.\n",
    "        hidden_state_size : int\n",
    "            Size of the hidden states (the input will be padded with 0's to this size).\n",
    "        message_size : int\n",
    "            Message function output vector size.\n",
    "        n_layers : int\n",
    "            Number of iterations Message+Update (weight tying).\n",
    "        l_target : int\n",
    "            Size of the output.\n",
    "        type : str (Optional)\n",
    "            Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, e, hidden_state_size, message_size, n_layers, l_target, type='regression'):\n",
    "        super(MpnnGGNN, self).__init__()\n",
    "\n",
    "        # Define message\n",
    "        self.m = nn.ModuleList([MessageFunction('ggnn', args={'e_label': e, 'in': hidden_state_size, 'out': message_size})])\n",
    "\n",
    "        # Define Update\n",
    "        self.u = nn.ModuleList([UpdateFunction('ggnn',\n",
    "                                                args={'in_m': message_size,\n",
    "                                                'out': hidden_state_size})])\n",
    "\n",
    "        # Define Readout\n",
    "        self.r = ReadoutFunction('ggnn',\n",
    "                                 args={'in': hidden_state_size,\n",
    "                                       'target': l_target})\n",
    "\n",
    "        self.type = type\n",
    "\n",
    "        self.args = {}\n",
    "        self.args['out'] = hidden_state_size\n",
    "\n",
    "        self.n_layers = n_layers\n",
    "\n",
    "    def forward(self, g, h_in, e):\n",
    "\n",
    "        h = []\n",
    "\n",
    "        # Padding to some larger dimension d\n",
    "        h_t = torch.cat([h_in, Variable(torch.Tensor(h_in.size(0), h_in.size(1), h_in.size(2)).type_as(h_in.data).zero_())], 2)\n",
    "        \n",
    "\n",
    "        h.append(h_t.clone())\n",
    "\n",
    "        # Layer\n",
    "        for t in range(0, self.n_layers):\n",
    "\n",
    "            h_t = Variable(torch.zeros(h[0].size(0), h[0].size(1), h[0].size(2)).type_as(h_in.data))\n",
    "\n",
    "            # Apply one layer pass (Message + Update)\n",
    "            for v in range(0, h_in.size(1)):\n",
    "                print(h[t].shape)\n",
    "                print(e.shape)\n",
    "                m = self.m[0].forward(h[t][:, v, :], h[t], e[:, v, :])\n",
    "                #m = self.m[0].forward(h[t][:, v, :], h[t], e[:, v, :])\n",
    "\n",
    "                # Nodes without edge set message to 0\n",
    "                m = g[:, v, :, None].expand_as(m) * m\n",
    "\n",
    "                m = torch.sum(m, 1)\n",
    "\n",
    "                # Update\n",
    "                h_t[:, v, :] = self.u[0].forward(h[t][:, v, :], m)\n",
    "\n",
    "            # Delete virtual nodes\n",
    "            h_t = (torch.sum(h_in, 2)[..., None].expand_as(h_t) > 0).type_as(h_t) * h_t\n",
    "            h.append(h_t.clone())\n",
    "\n",
    "        # Readout\n",
    "        res = self.r.forward(h)\n",
    "        if self.type == 'classification':\n",
    "            res = nn.LogSoftmax()(res)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MpnnDuvenaud(nn.Module):\n",
    "    \"\"\"\n",
    "        MPNN as proposed by Duvenaud et al..\n",
    "\n",
    "        This class implements the whole Duvenaud et al. model following the functions proposed by Gilmer et al. as \n",
    "        Message, Update and Readout.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        d : int list.\n",
    "            Possible degrees for the input graph.\n",
    "        in_n : int list\n",
    "            Sizes for the node and edge features.\n",
    "        out_update : int list\n",
    "            Output sizes for the different Update functions.\n",
    "        hidden_state_readout : int\n",
    "            Input size for the neural net used inside the readout function.\n",
    "        l_target : int\n",
    "            Size of the output.\n",
    "        type : str (Optional)\n",
    "            Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, d, in_n, out_update, hidden_state_readout, l_target, type='regression'):\n",
    "        super(MpnnDuvenaud, self).__init__()\n",
    "\n",
    "        n_layers = len(out_update)\n",
    "\n",
    "        # Define message 1 & 2\n",
    "        self.m = nn.ModuleList([MessageFunction('duvenaud') for _ in range(n_layers)])\n",
    "\n",
    "        # Define Update 1 & 2\n",
    "        self.u = nn.ModuleList([UpdateFunction('duvenaud', args={'deg': d, 'in': self.m[i].get_out_size(in_n[0], in_n[1]), 'out': out_update[0]}) if i == 0 else\n",
    "                                UpdateFunction('duvenaud', args={'deg': d, 'in': self.m[i].get_out_size(out_update[i-1], in_n[1]), 'out': out_update[i]}) for i in range(n_layers)])\n",
    "\n",
    "        # Define Readout\n",
    "        self.r = ReadoutFunction('duvenaud',\n",
    "                                 args={'layers': len(self.m) + 1,\n",
    "                                       'in': [in_n[0] if i == 0 else out_update[i-1] for i in range(n_layers+1)],\n",
    "                                       'out': hidden_state_readout,\n",
    "                                       'target': l_target})\n",
    "\n",
    "        self.type = type\n",
    "\n",
    "    def forward(self, g, h_in, e, plotter=None):\n",
    "\n",
    "        h = []\n",
    "        h.append(h_in)\n",
    "\n",
    "        # Layer\n",
    "        for t in range(0, len(self.m)):\n",
    "\n",
    "            u_args = self.u[t].get_args()\n",
    "\n",
    "            h_t = Variable(torch.zeros(h_in.size(0), h_in.size(1), u_args['out']).type_as(h[t].data))\n",
    "\n",
    "            # Apply one layer pass (Message + Update)\n",
    "            for v in range(0, h_in.size(1)):\n",
    "\n",
    "                m = self.m[t].forward(h[t][:, v, :], h[t], e[:, v, :])\n",
    "\n",
    "                # Nodes without edge set message to 0\n",
    "                m = g[:, v, :, None].expand_as(m) * m\n",
    "\n",
    "                m = torch.sum(m, 1)\n",
    "\n",
    "                # Duvenaud\n",
    "                deg = torch.sum(g[:, v, :].data, 1)\n",
    "\n",
    "                # Separate degrees\n",
    "                for i in range(len(u_args['deg'])):\n",
    "                    ind = deg == u_args['deg'][i]\n",
    "                    ind = Variable(torch.squeeze(torch.nonzero(torch.squeeze(ind))), volatile=True)\n",
    "\n",
    "                    opt = {'deg': i}\n",
    "\n",
    "                    # Update\n",
    "                    if len(ind) != 0:\n",
    "                        aux = self.u[t].forward(torch.index_select(h[t], 0, ind)[:, v, :], torch.index_select(m, 0, ind), opt)\n",
    "\n",
    "                        ind = ind.data.cpu().numpy()\n",
    "                        for j in range(len(ind)):\n",
    "                            h_t[ind[j], v, :] = aux[j, :]\n",
    "\n",
    "            if plotter is not None:\n",
    "                num_feat = h_t.size(2)\n",
    "                color = h_t[0,:,:].data.cpu().numpy()\n",
    "                for i in range(num_feat):\n",
    "                    plotter(color[:, i], 'layer_' + str(t) + '_element_' + str(i) + '.png')\n",
    "\n",
    "            h.append(h_t.clone())\n",
    "        # Readout\n",
    "        res = self.r.forward(h)\n",
    "        if self.type == 'classification':\n",
    "            res = nn.LogSoftmax()(res)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MPNN(nn.Module):\n",
    "    \"\"\"\n",
    "        MPNN as proposed by Gilmer et al..\n",
    "\n",
    "        This class implements the whole Gilmer et al. model following the functions Message, Update and Readout.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        in_n : int list\n",
    "            Sizes for the node and edge features.\n",
    "        hidden_state_size : int\n",
    "            Size of the hidden states (the input will be padded with 0's to this size).\n",
    "        message_size : int\n",
    "            Message function output vector size.\n",
    "        n_layers : int\n",
    "            Number of iterations Message+Update (weight tying).\n",
    "        l_target : int\n",
    "            Size of the output.\n",
    "        type : str (Optional)\n",
    "            Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_n, hidden_state_size, message_size, n_layers, l_target, type='regression'):\n",
    "        super(MPNN, self).__init__()\n",
    "\n",
    "        # Define message\n",
    "        self.m = nn.ModuleList(\n",
    "            [MessageFunction('mpnn', args={'edge_feat': in_n[1], 'in': hidden_state_size, 'out': message_size})])\n",
    "\n",
    "        # Define Update\n",
    "        self.u = nn.ModuleList([UpdateFunction('mpnn',\n",
    "                                               args={'in_m': message_size,\n",
    "                                                     'out': hidden_state_size})])\n",
    "\n",
    "        # Define Readout\n",
    "        self.r = ReadoutFunction('mpnn',\n",
    "                                 args={'in': hidden_state_size,\n",
    "                                       'target': l_target})\n",
    "\n",
    "        self.type = type\n",
    "\n",
    "        self.args = {}\n",
    "        self.args['out'] = hidden_state_size\n",
    "\n",
    "        self.n_layers = n_layers\n",
    "\n",
    "    def forward(self, g, h_in, e):\n",
    "\n",
    "        h = []\n",
    "\n",
    "        # Padding to some larger dimension d\n",
    "        h_t = torch.cat([h_in, Variable(\n",
    "            torch.zeros(h_in.size(0), h_in.size(1), self.args['out'] - h_in.size(2)).type_as(h_in.data))], 2)\n",
    "\n",
    "        h.append(h_t.clone())\n",
    "\n",
    "        # Layer\n",
    "        for t in range(0, self.n_layers):\n",
    "            e_aux = e.view(-1, e.size(3))\n",
    "\n",
    "            h_aux = h[t].view(-1, h[t].size(2))\n",
    "\n",
    "            m = self.m[0].forward(h[t], h_aux, e_aux)\n",
    "            m = m.view(h[0].size(0), h[0].size(1), -1, m.size(1))\n",
    "\n",
    "            # Nodes without edge set message to 0\n",
    "            m = torch.unsqueeze(g, 3).expand_as(m) * m\n",
    "\n",
    "            m = torch.squeeze(torch.sum(m, 1))\n",
    "\n",
    "            h_t = self.u[0].forward(h[t], m)\n",
    "\n",
    "            # Delete virtual nodes\n",
    "            h_t = (torch.sum(h_in, 2)[..., None].expand_as(h_t) > 0).type_as(h_t) * h_t\n",
    "            h.append(h_t)\n",
    "\n",
    "        # Readout\n",
    "        res = self.r.forward(h)\n",
    "\n",
    "        if self.type == 'classification':\n",
    "            res = nn.LogSoftmax()(res)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def qm9_nodes(g, hydrogen=False):\n",
    "    h = []\n",
    "    for n, d in g.nodes_iter(data=True):\n",
    "        h_t = []\n",
    "        # Atom type (One-hot H, C, N, O F)\n",
    "        h_t += [int(d['a_type'] == x) for x in ['H', 'C', 'N', 'O', 'F']]\n",
    "        # Atomic number\n",
    "        h_t.append(d['a_num'])\n",
    "        # Partial Charge\n",
    "        h_t.append(d['pc'])\n",
    "        # Acceptor\n",
    "        h_t.append(d['acceptor'])\n",
    "        # Donor\n",
    "        h_t.append(d['donor'])\n",
    "        # Aromatic\n",
    "        h_t.append(int(d['aromatic']))\n",
    "        # Hybradization\n",
    "        h_t += [int(d['hybridization'] == x) for x in [rdkit.Chem.rdchem.HybridizationType.SP, rdkit.Chem.rdchem.HybridizationType.SP2, rdkit.Chem.rdchem.HybridizationType.SP3]]\n",
    "        # If number hydrogen is used as a\n",
    "        if hydrogen:\n",
    "            h_t.append(d['num_h'])\n",
    "        h.append(h_t)\n",
    "    return h\n",
    "\n",
    "\n",
    "def qm9_edges(g, e_representation='raw_distance'):\n",
    "    remove_edges = []\n",
    "    e={}    \n",
    "    for n1, n2, d in g.edges_iter(data=True):\n",
    "        e_t = []\n",
    "        # Raw distance function\n",
    "        if e_representation == 'chem_graph':\n",
    "            if d['b_type'] is None:\n",
    "                remove_edges += [(n1, n2)]\n",
    "            else:\n",
    "                e_t += [i+1 for i, x in enumerate([rdkit.Chem.rdchem.BondType.SINGLE, rdkit.Chem.rdchem.BondType.DOUBLE,\n",
    "                                                rdkit.Chem.rdchem.BondType.TRIPLE, rdkit.Chem.rdchem.BondType.AROMATIC])\n",
    "                        if x == d['b_type']]\n",
    "        elif e_representation == 'distance_bin':\n",
    "            if d['b_type'] is None:\n",
    "                step = (6-2)/8.0\n",
    "                start = 2\n",
    "                b = 9\n",
    "                for i in range(0, 9):\n",
    "                    if d['distance'] < (start+i*step):\n",
    "                        b = i\n",
    "                        break\n",
    "                e_t.append(b+5)\n",
    "            else:\n",
    "                e_t += [i+1 for i, x in enumerate([rdkit.Chem.rdchem.BondType.SINGLE, rdkit.Chem.rdchem.BondType.DOUBLE,\n",
    "                                                   rdkit.Chem.rdchem.BondType.TRIPLE, rdkit.Chem.rdchem.BondType.AROMATIC])\n",
    "                        if x == d['b_type']]\n",
    "        elif e_representation == 'raw_distance':\n",
    "            if d['b_type'] is None:\n",
    "                remove_edges += [(n1, n2)]\n",
    "            else:\n",
    "                e_t.append(d['distance'])\n",
    "                e_t += [int(d['b_type'] == x) for x in [rdkit.Chem.rdchem.BondType.SINGLE, rdkit.Chem.rdchem.BondType.DOUBLE,\n",
    "                                                        rdkit.Chem.rdchem.BondType.TRIPLE, rdkit.Chem.rdchem.BondType.AROMATIC]]\n",
    "        else:\n",
    "            print('Incorrect Edge representation transform')\n",
    "            quit()\n",
    "        if e_t:\n",
    "            e[(n1, n2)] = e_t\n",
    "    for edg in remove_edges:\n",
    "        g.remove_edge(*edg)\n",
    "    return nx.to_numpy_matrix(g), e\n",
    "    \n",
    "\n",
    "def normalize_data(data, mean, std):\n",
    "    data_norm = (data-mean)/std\n",
    "    return data_norm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_graph_our(g_tag,prop):\n",
    "    \n",
    "    prop = prop.split()\n",
    "    g_H=float(prop[0])\n",
    "    g_p=float(prop[1])\n",
    "\n",
    "    labels = [g_H, g_p]\n",
    "    return nx.Graph(tag=g_tag,H=g_H, p=g_p), labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def xyz_graph_reader(graph_file):\n",
    "\n",
    "    with open(graph_file,'r') as f:\n",
    "        # Number of atoms\n",
    "        #na = int(f.readline())\n",
    "        tag=f.readline().strip()\n",
    "        # Graph properties\n",
    "        properties = f.readline()\n",
    "        na=int(f.readline())\n",
    "        #g, l = init_graph(properties)\n",
    "        g,l=init_graph_our(tag, properties)\n",
    "        atom_properties = []\n",
    "        # Atoms properties\n",
    "        for i in range(na):\n",
    "            a_properties = f.readline()\n",
    "            a_properties = a_properties.replace('.*^', 'e')\n",
    "            a_properties = a_properties.replace('*^', 'e')\n",
    "            a_properties = a_properties.split()\n",
    "            atom_properties.append(a_properties)\n",
    "\n",
    "        # Frequencies\n",
    "        #f.readline()\n",
    "\n",
    "        # SMILES\n",
    "        smiles = f.readline()\n",
    "        smiles = smiles.split()\n",
    "        smiles = smiles[0]\n",
    "        \n",
    "        m = Chem.MolFromSmiles(smiles)\n",
    "        m = Chem.AddHs(m)\n",
    "        \n",
    "        fdef_name = os.path.join(RDConfig.RDDataDir, 'BaseFeatures.fdef')\n",
    "        factory = ChemicalFeatures.BuildFeatureFactory(fdef_name)\n",
    "        feats = factory.GetFeaturesForMol(m)\n",
    "\n",
    "        # Create nodes\n",
    "        for i in range(0, m.GetNumAtoms()):\n",
    "            atom_i = m.GetAtomWithIdx(i)\n",
    "\n",
    "            g.add_node(i, a_type=atom_i.GetSymbol(), a_num=atom_i.GetAtomicNum(), acceptor=0, donor=0,\n",
    "                       aromatic=atom_i.GetIsAromatic(), hybridization=atom_i.GetHybridization(),\n",
    "                       num_h=atom_i.GetTotalNumHs(), coord=np.array(atom_properties[i][1:4]).astype(np.float),\n",
    "                       pc=float(atom_properties[i][4]))\n",
    "\n",
    "        for i in range(0, len(feats)):\n",
    "            if feats[i].GetFamily() == 'Donor':\n",
    "                node_list = feats[i].GetAtomIds()\n",
    "                for i in node_list:\n",
    "                    g.node[i]['donor'] = 1\n",
    "            elif feats[i].GetFamily() == 'Acceptor':\n",
    "                node_list = feats[i].GetAtomIds()\n",
    "                for i in node_list:\n",
    "                    g.node[i]['acceptor'] = 1\n",
    "\n",
    "        # Read Edges\n",
    "        for i in range(0, m.GetNumAtoms()):\n",
    "            for j in range(0, m.GetNumAtoms()):\n",
    "                e_ij = m.GetBondBetweenAtoms(i, j)\n",
    "                if e_ij is not None:\n",
    "                    g.add_edge(i, j, b_type=e_ij.GetBondType(),\n",
    "                               distance=np.linalg.norm(g.node[i]['coord']-g.node[j]['coord']))\n",
    "                else:\n",
    "                    # Unbonded\n",
    "                    g.add_edge(i, j, b_type=None,\n",
    "                               distance=np.linalg.norm(g.node[i]['coord'] - g.node[j]['coord']))\n",
    "    return g , l\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def restricted_float(x, inter):\n",
    "    x = float(x)\n",
    "    if x < inter[0] or x > inter[1]:\n",
    "        raise argparse.ArgumentTypeError(\"%r not in range [1e-5, 1e-4]\"%(x,))\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Qm9(data.Dataset):\n",
    "\n",
    "    # Constructor\n",
    "    def __init__(self, root_path, ids, vertex_transform=utils.qm9_nodes, edge_transform=utils.qm9_edges,\n",
    "                 target_transform=None, e_representation='raw_distance'):\n",
    "        self.root = root_path\n",
    "        self.ids = ids\n",
    "        self.vertex_transform = vertex_transform\n",
    "        self.edge_transform = edge_transform\n",
    "        self.target_transform = target_transform\n",
    "        self.e_representation = e_representation\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        g, target = xyz_graph_reader(os.path.join(self.root, self.ids[index]))\n",
    "        if self.vertex_transform is not None:\n",
    "            h = self.vertex_transform(g)\n",
    "\n",
    "        if self.edge_transform is not None:\n",
    "            g, e = utils.qm9_edges(g, self.e_representation)\n",
    "\n",
    "        if self.target_transform is not None:\n",
    "            target = self.target_transform(target)\n",
    "\n",
    "        return (g, h, e), target\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.ids)\n",
    "\n",
    "    def set_target_transform(self, target_transform):\n",
    "        self.target_transform = target_transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = args.root[0]\n",
    "files = [f for f in os.listdir(root) if os.path.isfile(os.path.join(root, f))]\n",
    "\n",
    "idx = np.random.permutation(len(files))\n",
    "idx = idx.tolist()\n",
    "\n",
    "valid_ids = [files[i] for i in idx[0:1000]]\n",
    "test_ids  = [files[i] for i in idx[1000:2000]]\n",
    "train_ids = [files[i] for i in idx[2000:]]\n",
    "\n",
    "data_train = Qm9(root, train_ids, vertex_transform=utils.qm9_nodes, edge_transform=lambda g: utils.qm9_edges(g, e_representation='raw_distance'))\n",
    "data_valid = Qm9(root, valid_ids)\n",
    "data_test = Qm9(root, test_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AverageMeter(object):\n",
    "    \"\"\"Computes and stores the average and current value\"\"\"\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.val = 0\n",
    "        self.avg = 0\n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "\n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val * n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.get_graph_stats(data_valid, 'degrees')\n",
    "stat_dict = utils.get_graph_stats(data_valid, ['target_mean', 'target_std'])\n",
    "\n",
    "data_train.set_target_transform(lambda x: utils.normalize_data(x,stat_dict['target_mean'],stat_dict['target_std']))\n",
    "data_valid.set_target_transform(lambda x: utils.normalize_data(x, stat_dict['target_mean'],stat_dict['target_std']))\n",
    "data_test.set_target_transform(lambda x: utils.normalize_data(x, stat_dict['target_mean'], stat_dict['target_std']))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = torch.utils.data.DataLoader(data_train,\n",
    "                                           batch_size=args.batch_size, shuffle=True,\n",
    "                                           collate_fn=utils.collate_g,\n",
    "                                        pin_memory=True)\n",
    "valid_loader = torch.utils.data.DataLoader(data_valid,\n",
    "                                           batch_size=args.batch_size, collate_fn=utils.collate_g,\n",
    "                                           pin_memory=True)\n",
    "test_loader = torch.utils.data.DataLoader(data_test,\n",
    "                                          batch_size=args.batch_size, collate_fn=utils.collate_g,\n",
    "                                          pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(train_loader, model, criterion, optimizer, epoch, evaluation, logger):\n",
    "    batch_time = AverageMeter()\n",
    "    data_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    error_ratio = AverageMeter()\n",
    "\n",
    "    # switch to train mode\n",
    "    model.train()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (g, h, e, target) in enumerate(train_loader):\n",
    "\n",
    "        # Prepare input data\n",
    "        if args.cuda:\n",
    "            g, h, e, target = g.cuda(), h.cuda(), e.cuda(), target.cuda()\n",
    "        g, h, e, target = Variable(g), Variable(h), Variable(e), Variable(target)\n",
    "\n",
    "        # Measure data loading time\n",
    "        data_time.update(time.time() - end)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Compute output\n",
    "        output = model(g, h, e)\n",
    "        train_loss = criterion(output, target)\n",
    "\n",
    "        # Logs\n",
    "        losses.update(train_loss.item(), g.size(0))\n",
    "        error_ratio.update(evaluation(output, target).item(), g.size(0))\n",
    "\n",
    "        # compute gradient and do SGD step\n",
    "        train_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # Measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % args.log_interval == 0 and i > 0:\n",
    "\n",
    "            print('Epoch: [{0}][{1}/{2}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Error Ratio {err.val:.4f} ({err.avg:.4f})'\n",
    "                  .format(epoch+1, i, len(train_loader), batch_time=batch_time,\n",
    "                          data_time=data_time, loss=losses, err=error_ratio))\n",
    "                          \n",
    "\n",
    "    print('Epoch: [{0}] Avg Error Ratio {err.avg:.3f}; Average Loss {loss.avg:.3f}; Avg Time x Batch {b_time.avg:.3f}'\n",
    "          .format(epoch, err=error_ratio, loss=losses, b_time=batch_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate(val_loader, model, criterion, evaluation, logger=None):\n",
    "    batch_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    error_ratio = AverageMeter()\n",
    "\n",
    "    # switch to evaluate mode\n",
    "    model.eval()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (g, h, e, target) in enumerate(val_loader):\n",
    "\n",
    "        # Prepare input data\n",
    "        if args.cuda:\n",
    "            g, h, e, target = g.cuda(), h.cuda(), e.cuda(), target.cuda()\n",
    "        g, h, e, target = Variable(g), Variable(h), Variable(e), Variable(target)\n",
    "\n",
    "        # Compute output\n",
    "        output = model(g, h, e)\n",
    "\n",
    "        # Logs\n",
    "        losses.update(criterion(output, target).item(), g.size(0))\n",
    "        error_ratio.update(evaluation(output, target).item(), g.size(0))\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % args.log_interval == 0 and i > 0:\n",
    "            \n",
    "            print('Test: [{0}/{1}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Error Ratio {err.val:.4f} ({err.avg:.4f})'\n",
    "                  .format(i, len(val_loader), batch_time=batch_time,\n",
    "                          loss=losses, err=error_ratio))\n",
    "\n",
    "    print(' * Average Error Ratio {err.avg:.3f}; Average Loss {loss.avg:.3f}'\n",
    "          .format(err=error_ratio, loss=losses))\n",
    "\n",
    "\n",
    "    return error_ratio.avg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "object of type 'int' has no len()",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-28-dcf244065566>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      7\u001b[0m \u001b[0ml_target\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0ml\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      8\u001b[0m \u001b[0mtype\u001b[0m \u001b[1;33m=\u001b[0m\u001b[1;34m'regression'\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 9\u001b[1;33m \u001b[0mmodel\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mMpnnDuvenaud\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mutils\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mget_graph_stats\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdata_valid\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'degrees'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m\u001b[0min_n\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mhidden_state_size\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmessage_size\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0ml_target\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtype\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mtype\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     10\u001b[0m \u001b[1;32mdel\u001b[0m \u001b[0min_n\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mhidden_state_size\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmessage_size\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mn_layers\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0ml_target\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtype\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     11\u001b[0m \u001b[0moptimizer\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0moptim\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mAdam\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mparameters\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlr\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlr\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-22-72a3e9aa0def>\u001b[0m in \u001b[0;36m__init__\u001b[1;34m(self, d, in_n, out_update, hidden_state_readout, l_target, type)\u001b[0m\n\u001b[0;32m     37\u001b[0m         \u001b[0msuper\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mMpnnDuvenaud\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m__init__\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     38\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 39\u001b[1;33m         \u001b[0mn_layers\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout_update\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     40\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     41\u001b[0m         \u001b[1;31m# Define message 1 & 2\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mTypeError\u001b[0m: object of type 'int' has no len()"
     ]
    }
   ],
   "source": [
    "g_tuple, l = data_train[0]\n",
    "g, h_t, e = g_tuple\n",
    "in_n = [len(h_t[0]), len(list(e.values())[0])]\n",
    "hidden_state_size = 73\n",
    "message_size = 73\n",
    "n_layers = 3\n",
    "l_target = len(l)\n",
    "type ='regression'\n",
    "model = MpnnDuvenaud(utils.get_graph_stats(data_valid, 'degrees'),in_n, hidden_state_size, message_size, l_target, type=type)\n",
    "del in_n, hidden_state_size, message_size, n_layers, l_target, type\n",
    "optimizer = optim.Adam(model.parameters(), lr=args.lr)\n",
    "criterion = nn.MSELoss()\n",
    "evaluation = lambda output, target: torch.mean(torch.abs(output - target) / torch.abs(target))\n",
    "lr_step = (args.lr-args.lr*args.lr_decay)/(args.epochs*args.schedule[1] - args.epochs*args.schedule[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=> loading best model './checkpoint/our_data/mpnn/model_best.pth'\n",
      "=> loaded best model './checkpoint/our_data/mpnn/model_best.pth' (epoch 44)\n"
     ]
    }
   ],
   "source": [
    "if args.resume:\n",
    "    checkpoint_dir = args.resume\n",
    "    best_model_file = os.path.join(checkpoint_dir, 'model_best.pth')\n",
    "    if not os.path.isdir(checkpoint_dir):\n",
    "        os.makedirs(checkpoint_dir)\n",
    "    if os.path.isfile(best_model_file):\n",
    "        print(\"=> loading best model '{}'\".format(best_model_file))\n",
    "        checkpoint = torch.load(best_model_file)\n",
    "        args.start_epoch = checkpoint['epoch']\n",
    "        best_acc1 = checkpoint['best_er1']\n",
    "        model.load_state_dict(checkpoint['state_dict'])\n",
    "        optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "        print(\"=> loaded best model '{}' (epoch {})\".format(best_model_file, checkpoint['epoch']))\n",
    "    else:\n",
    "        print(\"=> no best model found at '{}'\".format(best_model_file))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t* Cuda\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "MpnnGGNN(\n",
       "  (m): ModuleList(\n",
       "    (0): MessageFunction(\n",
       "      (learn_args): ParameterList(  (0): Parameter containing: [torch.cuda.FloatTensor of size 25x73x73 (GPU 0)])\n",
       "      (learn_modules): ModuleList()\n",
       "    )\n",
       "  )\n",
       "  (u): ModuleList(\n",
       "    (0): UpdateFunction(\n",
       "      (learn_args): ParameterList()\n",
       "      (learn_modules): ModuleList(\n",
       "        (0): GRU(73, 73)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (r): ReadoutFunction(\n",
       "    (learn_args): ParameterList()\n",
       "    (learn_modules): ModuleList(\n",
       "      (0): NNet(\n",
       "        (fcs): ModuleList(\n",
       "          (0): Linear(in_features=146, out_features=128, bias=True)\n",
       "          (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "          (2): Linear(in_features=256, out_features=128, bias=True)\n",
       "          (3): Linear(in_features=128, out_features=2, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (1): NNet(\n",
       "        (fcs): ModuleList(\n",
       "          (0): Linear(in_features=73, out_features=128, bias=True)\n",
       "          (1): Linear(in_features=128, out_features=256, bias=True)\n",
       "          (2): Linear(in_features=256, out_features=128, bias=True)\n",
       "          (3): Linear(in_features=128, out_features=2, bias=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if args.cuda:\n",
    "    print('\\t* Cuda')\n",
    "    model = model.cuda()\n",
    "    criterion = criterion.cuda()\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA out of memory. Tried to allocate 1.12 GiB (GPU 0; 4.00 GiB total capacity; 1.84 GiB already allocated; 115.79 MiB free; 905.70 MiB cached)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-64-30cc6d64d058>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      7\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      8\u001b[0m     \u001b[1;31m# train for one epoch\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 9\u001b[1;33m     \u001b[0mtrain\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtrain_loader\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mepoch\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mevaluation\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlogger\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mNone\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     10\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     11\u001b[0m     \u001b[1;31m# evaluate on test set\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-55-9e2d2e9bdc86>\u001b[0m in \u001b[0;36mtrain\u001b[1;34m(train_loader, model, criterion, optimizer, epoch, evaluation, logger)\u001b[0m\n\u001b[0;32m     22\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     23\u001b[0m         \u001b[1;31m# Compute output\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 24\u001b[1;33m         \u001b[0moutput\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mg\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mh\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0me\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     25\u001b[0m         \u001b[0mtrain_loss\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtarget\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     26\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Anaconda3\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    539\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    540\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 541\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    542\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    543\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-45-b24c7bd1c3b0>\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, g, h_in, e)\u001b[0m\n\u001b[0;32m     61\u001b[0m             \u001b[0mh_aux\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mh\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mt\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mview\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m-\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mh\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mt\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msize\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m2\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     62\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 63\u001b[1;33m             \u001b[0mm\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mm\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mh\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mt\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mh_aux\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0me_aux\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     64\u001b[0m             \u001b[0mm\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mm\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mview\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mh\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msize\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mh\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msize\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m-\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mm\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msize\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     65\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-42-1593b236ee87>\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, h_v, h_w, e_vw, args)\u001b[0m\n\u001b[0;32m     11\u001b[0m     \u001b[1;31m# Message from h_v to h_w through e_vw\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     12\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mh_v\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mh_w\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0me_vw\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mNone\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 13\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mm_function\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mh_v\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mh_w\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0me_vw\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     14\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     15\u001b[0m     \u001b[1;31m# Set a message function\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-42-1593b236ee87>\u001b[0m in \u001b[0;36mm_mpnn\u001b[1;34m(self, h_v, h_w, e_vw, opt)\u001b[0m\n\u001b[0;32m    139\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mm_mpnn\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mh_v\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mh_w\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0me_vw\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mopt\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;33m{\u001b[0m\u001b[1;33m}\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    140\u001b[0m         \u001b[1;31m# Matrices for each edge\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 141\u001b[1;33m         \u001b[0medge_output\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlearn_modules\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0me_vw\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    142\u001b[0m         \u001b[0medge_output\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0medge_output\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mview\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m-\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;34m'out'\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;34m'in'\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    143\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Anaconda3\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    539\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    540\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 541\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    542\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    543\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-41-53a2a7ec0c88>\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, x)\u001b[0m\n\u001b[0;32m     12\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mn_hlayers\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     13\u001b[0m             \u001b[0mx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mF\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrelu\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfcs\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mi\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 14\u001b[1;33m         \u001b[0mx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfcs\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;33m-\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     15\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     16\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Anaconda3\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    539\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    540\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 541\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    542\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    543\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Anaconda3\\lib\\site-packages\\torch\\nn\\modules\\linear.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, input)\u001b[0m\n\u001b[0;32m     85\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     86\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 87\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mF\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlinear\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     88\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     89\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mextra_repr\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Anaconda3\\lib\\site-packages\\torch\\nn\\functional.py\u001b[0m in \u001b[0;36mlinear\u001b[1;34m(input, weight, bias)\u001b[0m\n\u001b[0;32m   1368\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdim\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m==\u001b[0m \u001b[1;36m2\u001b[0m \u001b[1;32mand\u001b[0m \u001b[0mbias\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1369\u001b[0m         \u001b[1;31m# fused op is marginally faster\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1370\u001b[1;33m         \u001b[0mret\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0maddmm\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mbias\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mweight\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mt\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1371\u001b[0m     \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1372\u001b[0m         \u001b[0moutput\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmatmul\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mweight\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mt\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 1.12 GiB (GPU 0; 4.00 GiB total capacity; 1.84 GiB already allocated; 115.79 MiB free; 905.70 MiB cached)"
     ]
    }
   ],
   "source": [
    "for epoch in range(0, args.epochs):\n",
    "\n",
    "    if epoch > args.epochs * args.schedule[0] and epoch < args.epochs * args.schedule[1]:\n",
    "        args.lr -= lr_step\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = args.lr\n",
    "\n",
    "    # train for one epoch\n",
    "    train(train_loader, model, criterion, optimizer, epoch, evaluation, logger=None)\n",
    "\n",
    "    # evaluate on test set\n",
    "    er1 = validate(valid_loader, model, criterion, evaluation)\n",
    "\n",
    "    is_best = er1 > best_er1\n",
    "    best_er1 = min(er1, best_er1)\n",
    "    utils.save_checkpoint({'epoch': epoch + 1, 'state_dict': model.state_dict(), 'best_er1': best_er1,\n",
    "                           'optimizer': optimizer.state_dict(), }, is_best=is_best, directory=args.resume)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test: [20/500]\tTime 0.046 (0.052)\tLoss 0.1834 (0.3290)\tError Ratio 0.8179 (1.4637)\n",
      "Test: [40/500]\tTime 0.035 (0.056)\tLoss 0.2218 (0.3783)\tError Ratio 1.0000 (1.3206)\n",
      "Test: [60/500]\tTime 0.054 (0.057)\tLoss 0.1126 (0.9725)\tError Ratio 1.0472 (1.3199)\n",
      "Test: [80/500]\tTime 0.072 (0.056)\tLoss 0.1723 (1.0757)\tError Ratio 1.2968 (1.3300)\n",
      "Test: [100/500]\tTime 0.050 (0.056)\tLoss 0.3681 (0.9656)\tError Ratio 1.4976 (1.2820)\n",
      "Test: [120/500]\tTime 0.045 (0.056)\tLoss 0.4640 (0.8590)\tError Ratio 1.3237 (1.3108)\n",
      "Test: [140/500]\tTime 0.053 (0.056)\tLoss 0.2501 (0.9449)\tError Ratio 1.2683 (1.3011)\n",
      "Test: [160/500]\tTime 0.104 (0.057)\tLoss 1.9813 (0.8884)\tError Ratio 0.8331 (1.4272)\n",
      "Test: [180/500]\tTime 0.043 (0.057)\tLoss 0.1123 (0.8331)\tError Ratio 0.7974 (1.3982)\n",
      "Test: [200/500]\tTime 0.054 (0.056)\tLoss 0.4142 (0.8103)\tError Ratio 1.5463 (1.3763)\n",
      "Test: [220/500]\tTime 0.041 (0.056)\tLoss 0.2909 (0.7667)\tError Ratio 1.3218 (1.3967)\n",
      "Test: [240/500]\tTime 0.057 (0.056)\tLoss 0.4908 (0.7543)\tError Ratio 1.0000 (1.4337)\n",
      "Test: [260/500]\tTime 0.040 (0.056)\tLoss 0.6585 (0.7245)\tError Ratio 1.2175 (1.4353)\n",
      "Test: [280/500]\tTime 0.039 (0.056)\tLoss 0.4187 (0.7452)\tError Ratio 1.0000 (1.5191)\n",
      "Test: [300/500]\tTime 0.047 (0.055)\tLoss 0.2011 (0.7161)\tError Ratio 1.0226 (1.5532)\n",
      "Test: [320/500]\tTime 0.044 (0.055)\tLoss 0.2777 (0.6954)\tError Ratio 1.0002 (1.5275)\n",
      "Test: [340/500]\tTime 0.039 (0.055)\tLoss 0.1172 (0.8114)\tError Ratio 1.0093 (1.5217)\n",
      "Test: [360/500]\tTime 0.043 (0.055)\tLoss 0.4028 (0.8080)\tError Ratio 0.8304 (1.5078)\n",
      "Test: [380/500]\tTime 0.046 (0.055)\tLoss 0.4713 (0.7995)\tError Ratio 5.5029 (1.5124)\n",
      "Test: [400/500]\tTime 0.073 (0.055)\tLoss 8.6489 (0.7999)\tError Ratio 1.0001 (1.4939)\n",
      "Test: [420/500]\tTime 0.049 (0.054)\tLoss 0.2642 (0.7972)\tError Ratio 4.5652 (1.4814)\n",
      "Test: [440/500]\tTime 0.056 (0.054)\tLoss 0.0173 (0.7849)\tError Ratio 0.8390 (1.5407)\n",
      "Test: [460/500]\tTime 0.049 (0.054)\tLoss 0.2259 (0.8245)\tError Ratio 1.1144 (1.5350)\n",
      "Test: [480/500]\tTime 0.105 (0.055)\tLoss 0.2913 (0.8086)\tError Ratio 0.7469 (1.5218)\n",
      " * Average Error Ratio 1.512; Average Loss 0.817\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1.5121180399656295"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validate(test_loader, model, criterion, evaluation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((matrix([[0., 1., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [1., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 1., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 1., 0., 1., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 1., 0., 0., 0., 0., 0., 0.],\n",
       "          [1., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 1., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 1., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 1., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 1., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 1., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 1.],\n",
       "          [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.,\n",
       "           0., 0., 0., 0., 0., 0., 0., 0.]]),\n",
       "  [[0, 1, 0, 0, 0, 6, -0.18, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, -0.18, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, 0.18, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, 0.18, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, -0.15, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, -0.15, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 0, 1, 0, 0, 7, -0.15, 0, 0, 0, 0, 1, 0],\n",
       "   [0, 0, 1, 0, 0, 7, -0.15, 0, 0, 0, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, -0.15, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, -0.15, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, -0.15, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, -0.15, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, -0.15, 0, 0, 1, 0, 1, 0],\n",
       "   [0, 1, 0, 0, 0, 6, -0.15, 0, 0, 1, 0, 1, 0],\n",
       "   [1, 0, 0, 0, 0, 1, 0.15, 0, 0, 0, 0, 0, 0],\n",
       "   [1, 0, 0, 0, 0, 1, 0.15, 0, 0, 0, 0, 0, 0],\n",
       "   [1, 0, 0, 0, 0, 1, 0.15, 0, 0, 0, 0, 0, 0],\n",
       "   [1, 0, 0, 0, 0, 1, 0.15, 0, 0, 0, 0, 0, 0],\n",
       "   [1, 0, 0, 0, 0, 1, 0.15, 0, 0, 0, 0, 0, 0],\n",
       "   [1, 0, 0, 0, 0, 1, 0.15, 0, 0, 0, 0, 0, 0],\n",
       "   [1, 0, 0, 0, 0, 1, 0.15, 0, 0, 0, 0, 0, 0],\n",
       "   [1, 0, 0, 0, 0, 1, 0.15, 0, 0, 0, 0, 0, 0],\n",
       "   [1, 0, 0, 0, 0, 1, 0.15, 0, 0, 0, 0, 0, 0],\n",
       "   [1, 0, 0, 0, 0, 1, 0.15, 0, 0, 0, 0, 0, 0]],\n",
       "  {(0, 1): [1.2534656796258925, 0, 0, 0, 1],\n",
       "   (0, 5): [3.0361755219354496, 0, 0, 0, 1],\n",
       "   (0, 14): [2.66070847707899, 1, 0, 0, 0],\n",
       "   (1, 2): [2.1738618746369327, 0, 0, 0, 1],\n",
       "   (1, 15): [2.659367415382839, 1, 0, 0, 0],\n",
       "   (2, 3): [3.4403024649004337, 0, 0, 0, 1],\n",
       "   (2, 16): [2.1604445491611215, 1, 0, 0, 0],\n",
       "   (3, 4): [4.276896263647273, 0, 0, 0, 1],\n",
       "   (3, 6): [4.27654322555028, 1, 0, 0, 0],\n",
       "   (4, 5): [4.768194882762239, 0, 0, 0, 1],\n",
       "   (4, 17): [5.419578412570484, 1, 0, 0, 0],\n",
       "   (5, 18): [6.729531960693849, 1, 0, 0, 0],\n",
       "   (6, 7): [4.767567645036617, 0, 1, 0, 0],\n",
       "   (7, 8): [6.588300412853075, 1, 0, 0, 0],\n",
       "   (8, 9): [7.503280541869669, 0, 0, 0, 1],\n",
       "   (8, 13): [8.281099977659972, 0, 0, 0, 1],\n",
       "   (9, 10): [7.88237108743302, 0, 0, 0, 1],\n",
       "   (9, 19): [1.0861613185894627, 1, 0, 0, 0],\n",
       "   (10, 11): [7.502680668267843, 0, 0, 0, 1],\n",
       "   (10, 20): [1.086060642874052, 1, 0, 0, 0],\n",
       "   (11, 12): [8.28090223103739, 0, 0, 0, 1],\n",
       "   (11, 21): [1.0861098517185084, 1, 0, 0, 0],\n",
       "   (12, 13): [8.88310080996495, 0, 0, 0, 1],\n",
       "   (12, 22): [1.0860070073438755, 1, 0, 0, 0],\n",
       "   (13, 23): [1.0861842477222732, 1, 0, 0, 0]}),\n",
       " array([ 0.48303525, -0.63015803]))"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}