GNN_int_logloss_2Edges_ACSF_Set2Set.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import torch\n",
    "import pandas as pd\n",
    "from torch_geometric.data import Data,DataLoader\n",
    "from functions import *\n",
    "from pytorch_util import *\n",
    "from torch.optim import Adam\n",
    "from torch.nn.utils import clip_grad_value_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../Data/train_data_ACSF.pickle', 'rb') as handle:\n",
    "    train_data = pickle.load(handle)\n",
    "with open('../Data/val_data_ACSF.pickle', 'rb') as handle:\n",
    "    val_data = pickle.load(handle)\n",
    "with open('../Data/test_data_ACSF.pickle', 'rb') as handle:\n",
    "    test_data = pickle.load(handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "### parameters ###\n",
    "batch_size = 32\n",
    "dim = 64\n",
    "edge_dim = 12\n",
    "epochs = 5\n",
    "clip = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_list = [Data(**d) for d in train_data]\n",
    "train_dl = DataLoader(train_list,batch_size,shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_list = [Data(**d) for d in val_data]\n",
    "valid_dl = DataLoader(val_list,batch_size,shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_list = [Data(**d) for d in test_data]\n",
    "test_dl = DataLoader(test_list,batch_size,shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = next(iter(train_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net_int_2Edges_Set2Set3(torch.nn.Module):\n",
    "    # use set2set as head, pool->q_star, type->h0\n",
    "    def __init__(self,dim=64,edge_dim=12,node_in=8,edge_in=19,edge_in3=8,num_layers=1):\n",
    "        super(Net_int_2Edges_Set2Set3, self).__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.dim = dim\n",
    "        self.lin_node = torch.nn.Linear(node_in, dim)\n",
    "        self.lin_edge_attr = torch.nn.Linear(edge_in, edge_dim)\n",
    "        \n",
    "        nn1 = Linear(edge_dim, dim * dim, bias=False)\n",
    "        nn2 = Linear(edge_in3, dim * dim * 2 * 2, bias=False)\n",
    "        \n",
    "        self.conv1 = NNConv(dim, dim, nn1, aggr='mean', root_weight=False)\n",
    "        self.gru1 = GRU(dim, dim)\n",
    "        self.lin_covert = Sequential(BatchNorm1d(dim),Linear(dim, dim*2), \\\n",
    "                                     RReLU(), Dropout(),Linear(dim*2, dim*2),RReLU())\n",
    "        \n",
    "        self.conv2 = NNConv(dim*2, dim*2, nn2, aggr='mean', root_weight=False)\n",
    "        self.gru2 = GRU(dim*2, dim*2)\n",
    "        \n",
    "        self.lin_weight = Linear(edge_in3, dim*2*2, bias=False)\n",
    "        self.lin_bias = Linear(edge_in3, 1, bias=False)\n",
    "        \n",
    "        self.norm = BatchNorm1d(dim*2*2)\n",
    "        self.norm_x = BatchNorm1d(node_in)\n",
    "        self.head = Set2Set(dim*2,processing_steps=3,num_layers=num_layers)\n",
    "        self.pool = Set2Set(dim*2,processing_steps=3)\n",
    "        \n",
    "        self.h_lin = Linear(edge_in3,num_layers*dim*2*2)\n",
    "        self.q_star_lin = Linear(dim*2*2,dim*2*2)\n",
    "        \n",
    "    def forward(self, data,IsTrain=False):\n",
    "        out = F.rrelu(self.lin_node(self.norm_x(data.x)))\n",
    "        edge_attr = F.rrelu(self.lin_edge_attr(data.edge_attr))\n",
    "        h = out.unsqueeze(0)\n",
    "        # edge_*3 only does not repeat for undirected graph. Hence need to add (j,i) to (i,j) in edges\n",
    "        edge_index3 = torch.cat([data.edge_index3,data.edge_index3[[1,0]]],1)\n",
    "        edge_attr3 = torch.cat([data.edge_attr3,data.edge_attr3],0)\n",
    "        \n",
    "        for i in range(2):\n",
    "            # using bonding as edge\n",
    "            m = F.rrelu(self.conv1(out, data.edge_index, edge_attr))\n",
    "            out, h = self.gru1(m.unsqueeze(0), h)\n",
    "            out = out.squeeze(0)\n",
    "        \n",
    "        out = self.lin_covert(out)\n",
    "        h = out.unsqueeze(0)\n",
    "        for i in range(2):\n",
    "            # using couping as edge\n",
    "            m = F.rrelu(self.conv2(out, edge_index3, edge_attr3))\n",
    "            out, h = self.gru2(m.unsqueeze(0), h)\n",
    "            out = out.squeeze(0)  \n",
    "            \n",
    "        n = data.edge_index3.shape[1]\n",
    "        range_ = torch.arange(n)\n",
    "        batch_index = torch.cat([range_,range_]).to('cuda:0')\n",
    "        temp = out[data.edge_index3].reshape(2*n,-1) # (2*n_target,d)\n",
    "\n",
    "        coupling_batch_index = data.batch[data.edge_index3[0]]\n",
    "        pool = self.pool(out, data.batch) # (m,d)\n",
    "        pool = pool[coupling_batch_index] # (n_target,d)        \n",
    "        \n",
    "        h0,h1 = self.h_lin(data.edge_attr3).split(self.num_layers*self.dim*2,1)\n",
    "        h = (h0.reshape(-1,self.num_layers,self.dim*2).transpose(0,1).contiguous(),\\\n",
    "             h1.reshape(-1,self.num_layers,self.dim*2).transpose(0,1).contiguous())\n",
    "        q_star = self.q_star_lin(pool)\n",
    "            \n",
    "        yhat = self.head(temp,batch_index,h,q_star)  \n",
    "        yhat = self.norm(yhat)\n",
    "        \n",
    "        weight = self.lin_weight(data.edge_attr3)\n",
    "        bias = self.lin_bias(data.edge_attr3)\n",
    "        yhat = torch.sum(yhat * weight,1,keepdim=True) + bias\n",
    "        yhat = yhat.squeeze(1)\n",
    "        \n",
    "        if IsTrain:\n",
    "            k = torch.sum(data.edge_attr3,0)\n",
    "            nonzeroIndex = torch.nonzero(k).squeeze(1)\n",
    "            abs_ = torch.abs(data.y-yhat).unsqueeze(1)\n",
    "            loss = torch.sum(torch.log(torch.sum(abs_ * data.edge_attr3[:,nonzeroIndex],0)/k[nonzeroIndex]))/nonzeroIndex.shape[0]\n",
    "            return loss\n",
    "        else:\n",
    "            return yhat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### build model and set up training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = Net_int_2Edges_Set2Set3(dim=dim,edge_dim=edge_dim,node_in=89,num_layers=2).to('cuda:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "paras = trainable_parameter(model)\n",
    "opt = Adam(paras,lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:0, train_loss:1.047349210465723, val_loss:0.5481739627499865\n",
      "epoch:1, train_loss:0.5036085564821999, val_loss:0.3091499479726339\n",
      "epoch:2, train_loss:0.36555841940556916, val_loss:0.24013803814911944\n",
      "epoch:3, train_loss:0.2825971062448017, val_loss:0.1509034380189374\n",
      "epoch:4, train_loss:0.2217576992700143, val_loss:0.07818713443926893\n",
      "Training completed in 337.7606246471405s\n"
     ]
    }
   ],
   "source": [
    "# Net_int_2Edges_Set2Set3 (3 steps, num_layers=2)\n",
    "since = time.time()\n",
    "opt.zero_grad()\n",
    "for epoch in range(epochs):\n",
    "    # training #\n",
    "    model.train()\n",
    "    np.random.seed()\n",
    "    train_loss = 0\n",
    "    val_loss = 0\n",
    "    \n",
    "    for i,data in enumerate(train_dl):\n",
    "        data = data.to('cuda:0')\n",
    "        loss = model(data,True)\n",
    "        loss.backward()\n",
    "        clip_grad_value_(paras,clip)\n",
    "        opt.step()\n",
    "        opt.zero_grad()\n",
    "        train_loss += loss.item()\n",
    "\n",
    "    # evaluating #\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for j,data in enumerate(valid_dl):\n",
    "            data = data.to('cuda:0')\n",
    "            loss = model(data,True)\n",
    "            val_loss += loss.item()\n",
    "    print('epoch:{}, train_loss:{}, val_loss:{}'.format(epoch,train_loss/i,val_loss/j))\n",
    "    \n",
    "time_elapsed = time.time() - since\n",
    "print('Training completed in {}s'.format(time_elapsed))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:0, train_loss:1.0560779975983559, val_loss:0.5833897830711471\n",
      "epoch:1, train_loss:0.5029618185188981, val_loss:0.30379895493388176\n",
      "epoch:2, train_loss:0.3669695601746215, val_loss:0.21524867827757302\n",
      "epoch:3, train_loss:0.2803956470501221, val_loss:0.1342123432451079\n",
      "epoch:4, train_loss:0.2121753142776149, val_loss:0.07976430511245361\n",
      "Training completed in 258.1577181816101s\n"
     ]
    }
   ],
   "source": [
    "# Net_int_2Edges_Set2Set2 (3 steps, num_layers=2)\n",
    "since = time.time()\n",
    "opt.zero_grad()\n",
    "for epoch in range(epochs):\n",
    "    # training #\n",
    "    model.train()\n",
    "    np.random.seed()\n",
    "    train_loss = 0\n",
    "    val_loss = 0\n",
    "    \n",
    "    for i,data in enumerate(train_dl):\n",
    "        data = data.to('cuda:0')\n",
    "        loss = model(data,True)\n",
    "        loss.backward()\n",
    "        clip_grad_value_(paras,clip)\n",
    "        opt.step()\n",
    "        opt.zero_grad()\n",
    "        train_loss += loss.item()\n",
    "\n",
    "    # evaluating #\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for j,data in enumerate(valid_dl):\n",
    "            data = data.to('cuda:0')\n",
    "            loss = model(data,True)\n",
    "            val_loss += loss.item()\n",
    "    print('epoch:{}, train_loss:{}, val_loss:{}'.format(epoch,train_loss/i,val_loss/j))\n",
    "    \n",
    "time_elapsed = time.time() - since\n",
    "print('Training completed in {}s'.format(time_elapsed))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "# checkpoint = torch.load('../Model/gnn_int_logloss_2Edges.tar')\n",
    "# model.load_state_dict(checkpoint['model_state_dict'])\n",
    "# opt.load_state_dict(checkpoint['optimizer_state_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save({'model_state_dict': model.state_dict(),\n",
    "            'optimizer_state_dict': opt.state_dict(),\n",
    "            }, '../Model/Net_int_2Edges_pool_intHead2.tar')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### make submissions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "yhat_list = []\n",
    "with torch.no_grad():\n",
    "    for data in test_dl:\n",
    "        data = data.to('cuda:0')\n",
    "        yhat_list.append(model(data,False))\n",
    "\n",
    "yhat = torch.cat(yhat_list).cpu().detach().numpy()\n",
    "submission = pd.read_csv('../Data/sample_submission.csv')\n",
    "submission['scalar_coupling_constant'] = yhat\n",
    "submission.to_csv('../Submission/Net_int_2Edges_pool_intHead2.csv',index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}