notebooks/04_sudoku_train-4264619e.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "t_gnhLD0HvhU",
    "toc-hr-collapsed": true
   },
   "source": [
    "## Summary\n",
    "\n",
    "- *hidden_size = 162*.\n",
    "- *num_heads = 9*.\n",
    "- *dropout = 0*.\n",
    "- N=16.\n",
    "- Add node and edge features (node features as 81-dim. embedding in `hidden_size`-dim space).\n",
    "- Edgeconv: embed x and edge to half their size and keep row x only.\n",
    "- Embed attention with `model_size == 63` and add `output_dim` attribute to attention.\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install dependencies (Google Colab only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import google.colab\n",
    "\n",
    "    GOOGLE_COLAB = True\n",
    "except ImportError:\n",
    "    GOOGLE_COLAB = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 170
    },
    "colab_type": "code",
    "id": "hIbvKDk_HvhX",
    "outputId": "349cb6d3-ffbc-4550-d9dd-dd6288ed5356"
   },
   "outputs": [],
   "source": [
    "if GOOGLE_COLAB:\n",
    "    !pip install --upgrade torch-scatter\n",
    "    !pip install --upgrade torch-sparse\n",
    "    !pip install --upgrade torch-cluster\n",
    "    !pip install --upgrade torch-spline-conv\n",
    "    !pip install torch-geometric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 203
    },
    "colab_type": "code",
    "id": "tw-y9MmIHvha",
    "outputId": "0666864a-5017-40a1-a77c-d23388eb0b1c"
   },
   "outputs": [],
   "source": [
    "if GOOGLE_COLAB:\n",
    "    !pip install git+https://gitlab.com/ostrokach/proteinsolver.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "eoBMUoW2Hvhp",
    "toc-hr-collapsed": true
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TbKxMUZWHvhq"
   },
   "outputs": [],
   "source": [
    "import atexit\n",
    "import csv\n",
    "import itertools\n",
    "import tempfile\n",
    "import time\n",
    "import uuid\n",
    "import warnings\n",
    "from collections import deque\n",
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyarrow\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.utils.tensorboard\n",
    "from torch import optim\n",
    "from torch_geometric.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 334
    },
    "colab_type": "code",
    "id": "rczy7pPiHvhs",
    "outputId": "8e4673d7-4cc7-4f8e-ffc4-de94302e5fab"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/dev/shm/env/lib/python3.7/site-packages/Bio/KDTree/__init__.py:25: BiopythonDeprecationWarning: Bio.KDTree has been deprecated, and we intend to remove it in a future release of Biopython. Please use Bio.PDB.kdtrees instead, which is functionally very similar.\n",
      "  BiopythonDeprecationWarning)\n"
     ]
    }
   ],
   "source": [
    "import proteinsolver\n",
    "import proteinsolver.datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4fFnAyUOHvhv"
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/localscratch/strokach.3687035.0/sudoku')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DATA_ROOT = Path(tempfile.gettempdir())\n",
    "DATA_ROOT = next(Path(\"/localscratch/\").glob(\"strokach.*\")).joinpath(\"sudoku\")\n",
    "DATA_ROOT.mkdir(exist_ok=True)\n",
    "DATA_ROOT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "UNIQUE_ID = \"4264619e\"\n",
    "CONTINUE_PREVIOUS = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(PosixPath('/lustre04/scratch/strokach/workspace/proteinsolver/notebooks/sudoku_train'),\n",
       " PosixPath('/lustre04/scratch/strokach/workspace/proteinsolver/notebooks/sudoku_train/4264619e'))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "try:\n",
    "    NOTEBOOK_PATH\n",
    "    UNIQUE_PATH\n",
    "except NameError:\n",
    "    NOTEBOOK_PATH = Path(\"sudoku_train\").resolve()\n",
    "    NOTEBOOK_PATH.mkdir(exist_ok=True)\n",
    "    if UNIQUE_ID is None:\n",
    "        UNIQUE_ID = uuid.uuid4().hex[:8]\n",
    "        exist_ok = False\n",
    "    else:\n",
    "        exist_ok = True\n",
    "    UNIQUE_PATH = NOTEBOOK_PATH.joinpath(UNIQUE_ID)\n",
    "    UNIQUE_PATH.mkdir(exist_ok=exist_ok)\n",
    "NOTEBOOK_PATH, UNIQUE_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/scratch/strokach/datapkg_output_dir')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DATAPKG_DATA_DIR = Path(f\"~/datapkg_output_dir\").expanduser().resolve()\n",
    "DATAPKG_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/scratch/strokach/datapkg_output_dir'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "proteinsolver.settings.data_url = DATAPKG_DATA_DIR.as_posix()\n",
    "proteinsolver.settings.data_url"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `SudokuDataset`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZP3bLDO9PgT9"
   },
   "outputs": [],
   "source": [
    "for i in range(10):\n",
    "    dataset_name = f\"sudoku_train_{i}\"\n",
    "    datasets[dataset_name] = proteinsolver.datasets.SudokuDataset4(\n",
    "        root=DATA_ROOT.joinpath(dataset_name), subset=f\"train_{i}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets[\"sudoku_valid_0\"] = proteinsolver.datasets.SudokuDataset4(\n",
    "    root=DATA_ROOT.joinpath(\"sudoku_valid_0\"), subset=f\"valid_0\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "datasets[\"sudoku_valid_old\"] = proteinsolver.datasets.SudokuDataset2(\n",
    "    root=DATA_ROOT.joinpath(\"sudoku_valid_old\"),\n",
    "    data_url=DATAPKG_DATA_DIR.joinpath(\n",
    "        \"deep-protein-gen\", \"sudoku\", \"sudoku_valid.csv.gz\"\n",
    "    ).as_posix(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e3kWuRaxr89h"
   },
   "source": [
    "# Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tNCaSlrFDGA6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting /lustre04/scratch/strokach/workspace/proteinsolver/notebooks/sudoku_train/4264619e/model.py\n"
     ]
    }
   ],
   "source": [
    "%%file {UNIQUE_PATH}/model.py\n",
    "import copy\n",
    "import tempfile\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn.modules.container import ModuleList\n",
    "from torch_geometric.nn.inits import reset\n",
    "from torch_geometric.utils import (\n",
    "    add_self_loops,\n",
    "    remove_self_loops,\n",
    "    scatter_,\n",
    "    to_dense_adj,\n",
    "    to_dense_batch,\n",
    ")\n",
    "\n",
    "\n",
    "class EdgeConvMod(torch.nn.Module):\n",
    "    def __init__(self, nn, aggr=\"max\"):\n",
    "        super().__init__()\n",
    "        self.nn = nn\n",
    "        self.aggr = aggr\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        reset(self.nn)\n",
    "\n",
    "    def forward(self, x, edge_index, edge_attr=None):\n",
    "        \"\"\"\"\"\"\n",
    "        row, col = edge_index\n",
    "        x = x.unsqueeze(-1) if x.dim() == 1 else x\n",
    "\n",
    "        # TODO: Try -x[col] instead of x[col] - x[row]\n",
    "        if edge_attr is None:\n",
    "            out = torch.cat([x[row], x[col]], dim=-1)\n",
    "        else:\n",
    "            out = torch.cat([x[row], x[col], edge_attr], dim=-1)\n",
    "        out = self.nn(out)\n",
    "        x = scatter_(self.aggr, out, row, dim_size=x.size(0))\n",
    "\n",
    "        return x, out\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"{}(nn={})\".format(self.__class__.__name__, self.nn)\n",
    "\n",
    "\n",
    "class EdgeConvBatch(nn.Module):\n",
    "    def __init__(self, gnn, hidden_size, batch_norm=True, dropout=0.2):\n",
    "        super().__init__()\n",
    "\n",
    "        self.gnn = gnn\n",
    "\n",
    "        x_post_modules = []\n",
    "        edge_attr_post_modules = []\n",
    "\n",
    "        if batch_norm is not None:\n",
    "            x_post_modules.append(nn.LayerNorm(hidden_size))\n",
    "            edge_attr_post_modules.append(nn.LayerNorm(hidden_size))\n",
    "\n",
    "        if dropout:\n",
    "            x_post_modules.append(nn.Dropout(dropout))\n",
    "            edge_attr_post_modules.append(nn.Dropout(dropout))\n",
    "\n",
    "        self.x_postprocess = nn.Sequential(*x_post_modules)\n",
    "        self.edge_attr_postprocess = nn.Sequential(*edge_attr_post_modules)\n",
    "\n",
    "    def forward(self, x, edge_index, edge_attr=None):\n",
    "        x, edge_attr = self.gnn(x, edge_index, edge_attr)\n",
    "        x = self.x_postprocess(x)\n",
    "        edge_attr = self.edge_attr_postprocess(edge_attr)\n",
    "        return x, edge_attr\n",
    "\n",
    "\n",
    "def get_graph_conv_layer(input_size, hidden_size, output_size):\n",
    "    mlp = nn.Sequential(\n",
    "        #\n",
    "        nn.Linear(input_size, hidden_size),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(hidden_size, output_size),\n",
    "    )\n",
    "    gnn = EdgeConvMod(nn=mlp, aggr=\"add\")\n",
    "    graph_conv = EdgeConvBatch(gnn, output_size, batch_norm=True, dropout=0.2)\n",
    "    return graph_conv\n",
    "\n",
    "\n",
    "class MyEdgeConv(torch.nn.Module):\n",
    "    def __init__(self, hidden_size):\n",
    "        super().__init__()\n",
    "        self.embed_x = nn.Linear(hidden_size, hidden_size // 2)\n",
    "        self.embed_edge = nn.Linear(hidden_size, hidden_size // 2)\n",
    "\n",
    "        self.nn = nn.Sequential(\n",
    "            #\n",
    "            nn.Linear(hidden_size, hidden_size * 2),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size * 2, hidden_size),\n",
    "        )\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        reset(self.nn)\n",
    "\n",
    "    def forward(self, x, edge_index, edge_attr=None):\n",
    "        \"\"\"\"\"\"\n",
    "        row, col = edge_index\n",
    "        x = x.unsqueeze(-1) if x.dim() == 1 else x\n",
    "\n",
    "        x_in = self.embed_x(x)\n",
    "        edge_attr_in = self.embed_edge(edge_attr)\n",
    "        x_edge_attr_in = torch.cat([x_in[row], edge_attr_in], dim=-1)\n",
    "        edge_attr_out = self.nn(x_edge_attr_in)\n",
    "\n",
    "        #         if edge_attr is None:\n",
    "        #             out = torch.cat([x[row], x[col]], dim=-1)\n",
    "        #         else:\n",
    "        #             out = torch.cat([x[row], x[col], edge_attr], dim=-1)\n",
    "        #         edge_attr_out = self.nn(out)\n",
    "\n",
    "        return edge_attr_out\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"{}(nn={})\".format(self.__class__.__name__, self.nn)\n",
    "\n",
    "\n",
    "class MyAttn(torch.nn.Module):\n",
    "    def __init__(self, hidden_size):\n",
    "        super().__init__()\n",
    "        model_size = 63\n",
    "        self.embed_x = nn.Linear(hidden_size, model_size)\n",
    "        self.attn = MultiheadAttention(\n",
    "            embed_dim=model_size,\n",
    "            output_dim=hidden_size,\n",
    "            kdim=hidden_size,\n",
    "            vdim=hidden_size,\n",
    "            num_heads=9,\n",
    "            dropout=0,\n",
    "            bias=True,\n",
    "        )\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        reset(self.attn)\n",
    "\n",
    "    def forward(self, x, edge_index, edge_attr, batch):\n",
    "        \"\"\"\"\"\"\n",
    "        query = self.embed_x(x).unsqueeze(0)\n",
    "        key = to_dense_adj(edge_index, batch=batch, edge_attr=edge_attr).squeeze(0)\n",
    "\n",
    "        adjacency = to_dense_adj(edge_index, batch=batch).squeeze(0)\n",
    "        key_padding_mask = adjacency == 0\n",
    "        key_padding_mask[torch.eye(key_padding_mask.size(0)).to(torch.bool)] = 0\n",
    "        #         attn_mask = torch.zeros_like(key)\n",
    "        #         attn_mask[mask] = -float(\"inf\")\n",
    "\n",
    "        x_out, _ = self.attn(query, key, key, key_padding_mask=key_padding_mask)\n",
    "        #         x_out = torch.where(torch.isnan(x_out), torch.zeros_like(x_out), x_out)\n",
    "        x_out = x_out.squeeze(0)\n",
    "        assert (x_out == x_out).all().item()\n",
    "        assert x.shape == x_out.shape, (x.shape, x_out.shape)\n",
    "        return x_out\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"{}(nn={})\".format(self.__class__.__name__, self.nn)\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(\n",
    "        self, x_input_size, adj_input_size, hidden_size, output_size, batch_size=1\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        x_labels = torch.arange(81, dtype=torch.long)\n",
    "        self.register_buffer(\"x_labels\", x_labels)\n",
    "\n",
    "        self.register_buffer(\"batch\", torch.zeros(10000, dtype=torch.int64))\n",
    "\n",
    "        self.embed_x = nn.Sequential(nn.Embedding(x_input_size, hidden_size), nn.ReLU())\n",
    "        self.embed_x_labels = nn.Sequential(nn.Embedding(81, hidden_size), nn.ReLU())\n",
    "        self.finalize_x = nn.Sequential(\n",
    "            nn.Linear(hidden_size * 2, hidden_size), nn.LayerNorm(hidden_size)\n",
    "        )\n",
    "\n",
    "        if adj_input_size:\n",
    "            self.embed_adj = nn.Sequential(\n",
    "                nn.Linear(adj_input_size, hidden_size),\n",
    "                nn.ReLU(),\n",
    "                nn.Linear(hidden_size, hidden_size),\n",
    "                nn.LayerNorm(hidden_size),\n",
    "                #                 nn.ELU(),\n",
    "            )\n",
    "        else:\n",
    "            self.embed_adj = None\n",
    "\n",
    "        N = 16\n",
    "        self.N = N\n",
    "\n",
    "        norm = nn.LayerNorm(hidden_size)\n",
    "        self.x_norms_0 = _get_clones(norm, N)\n",
    "        self.adj_norms_0 = _get_clones(norm, N)\n",
    "        self.x_norms_1 = _get_clones(norm, N)\n",
    "        self.adj_norms_1 = _get_clones(norm, N)\n",
    "\n",
    "        edge_conv = MyEdgeConv(hidden_size)\n",
    "        self.edge_convs = _get_clones(edge_conv, N)\n",
    "\n",
    "        attn = MyAttn(hidden_size)\n",
    "        self.attns = _get_clones(attn, N)\n",
    "\n",
    "        self.dropout = nn.Dropout(0.1)\n",
    "\n",
    "        self.linear_out = nn.Linear(hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x, edge_index, edge_attr):\n",
    "\n",
    "        x = self.embed_x(x)\n",
    "        x_labels = self.embed_x_labels(self.x_labels)\n",
    "        x_labels = x_labels.repeat(x.size(0) // x_labels.size(0), 1)\n",
    "        x = torch.cat([x, x_labels], dim=1)\n",
    "        x = self.finalize_x(x)\n",
    "\n",
    "        edge_attr = self.embed_adj(edge_attr)\n",
    "\n",
    "        for i in range(self.N):\n",
    "            edge_attr_out = self.edge_convs[i](x, edge_index, edge_attr)\n",
    "            edge_attr = edge_attr + self.dropout(edge_attr_out)\n",
    "            edge_attr = self.adj_norms_1[i](edge_attr)\n",
    "\n",
    "            x_out = self.attns[i](\n",
    "                x,\n",
    "                edge_index,\n",
    "                self.adj_norms_0[i](edge_attr_out),\n",
    "                self.batch[: x.size(0)],\n",
    "            )\n",
    "            x = x + self.dropout(x_out)\n",
    "            x = self.x_norms_1[i](x)\n",
    "\n",
    "        x = self.linear_out(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class MultiheadAttention(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        embed_dim,\n",
    "        output_dim,\n",
    "        num_heads,\n",
    "        dropout=0.0,\n",
    "        bias=True,\n",
    "        add_bias_kv=False,\n",
    "        add_zero_attn=False,\n",
    "        kdim=None,\n",
    "        vdim=None,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.embed_dim = embed_dim\n",
    "        self.kdim = kdim if kdim is not None else embed_dim\n",
    "        self.vdim = vdim if vdim is not None else embed_dim\n",
    "        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim\n",
    "\n",
    "        self.num_heads = num_heads\n",
    "        self.dropout = dropout\n",
    "        self.head_dim = embed_dim // num_heads\n",
    "        assert (\n",
    "            self.head_dim * num_heads == self.embed_dim\n",
    "        ), \"embed_dim must be divisible by num_heads\"\n",
    "\n",
    "        if self._qkv_same_embed_dim is False:\n",
    "            self.q_proj_weight = nn.Parameter(torch.Tensor(embed_dim, embed_dim))\n",
    "            self.k_proj_weight = nn.Parameter(torch.Tensor(embed_dim, self.kdim))\n",
    "            self.v_proj_weight = nn.Parameter(torch.Tensor(embed_dim, self.vdim))\n",
    "        else:\n",
    "            self.in_proj_weight = nn.Parameter(torch.empty(3 * embed_dim, embed_dim))\n",
    "\n",
    "        if bias:\n",
    "            self.in_proj_bias = nn.Parameter(torch.empty(3 * embed_dim))\n",
    "        else:\n",
    "            self.register_parameter(\"in_proj_bias\", None)\n",
    "        self.out_proj = nn.Linear(embed_dim, output_dim, bias=bias)\n",
    "\n",
    "        if add_bias_kv:\n",
    "            self.bias_k = nn.Parameter(torch.empty(1, 1, embed_dim))\n",
    "            self.bias_v = nn.Parameter(torch.empty(1, 1, embed_dim))\n",
    "        else:\n",
    "            self.bias_k = self.bias_v = None\n",
    "\n",
    "        self.add_zero_attn = add_zero_attn\n",
    "\n",
    "        self._reset_parameters()\n",
    "\n",
    "    def _reset_parameters(self):\n",
    "        if self._qkv_same_embed_dim:\n",
    "            nn.init.xavier_uniform_(self.in_proj_weight)\n",
    "        else:\n",
    "            nn.init.xavier_uniform_(self.q_proj_weight)\n",
    "            nn.init.xavier_uniform_(self.k_proj_weight)\n",
    "            nn.init.xavier_uniform_(self.v_proj_weight)\n",
    "\n",
    "        if self.in_proj_bias is not None:\n",
    "            nn.init.constant_(self.in_proj_bias, 0.0)\n",
    "            nn.init.constant_(self.out_proj.bias, 0.0)\n",
    "        if self.bias_k is not None:\n",
    "            nn.init.xavier_normal_(self.bias_k)\n",
    "        if self.bias_v is not None:\n",
    "            nn.init.xavier_normal_(self.bias_v)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        query,\n",
    "        key,\n",
    "        value,\n",
    "        key_padding_mask=None,\n",
    "        need_weights=True,\n",
    "        attn_mask=None,\n",
    "    ):\n",
    "        if hasattr(self, \"_qkv_same_embed_dim\") and self._qkv_same_embed_dim is False:\n",
    "            return F.multi_head_attention_forward(\n",
    "                query,\n",
    "                key,\n",
    "                value,\n",
    "                self.embed_dim,\n",
    "                self.num_heads,\n",
    "                None,  # set self.in_proj_weight = None\n",
    "                self.in_proj_bias,\n",
    "                self.bias_k,\n",
    "                self.bias_v,\n",
    "                self.add_zero_attn,\n",
    "                self.dropout,\n",
    "                self.out_proj.weight,\n",
    "                self.out_proj.bias,\n",
    "                training=self.training,\n",
    "                key_padding_mask=key_padding_mask,\n",
    "                need_weights=need_weights,\n",
    "                attn_mask=attn_mask,\n",
    "                use_separate_proj_weight=True,\n",
    "                q_proj_weight=self.q_proj_weight,\n",
    "                k_proj_weight=self.k_proj_weight,\n",
    "                v_proj_weight=self.v_proj_weight,\n",
    "            )\n",
    "        else:\n",
    "            if not hasattr(self, \"_qkv_same_embed_dim\"):\n",
    "                warnings.warn(\n",
    "                    \"A new version of MultiheadAttention module has been implemented. \\\n",
    "                    Please re-train your model with the new module\",\n",
    "                    UserWarning,\n",
    "                )\n",
    "\n",
    "            return F.multi_head_attention_forward(\n",
    "                query,\n",
    "                key,\n",
    "                value,\n",
    "                self.embed_dim,\n",
    "                self.num_heads,\n",
    "                self.in_proj_weight,\n",
    "                self.in_proj_bias,\n",
    "                self.bias_k,\n",
    "                self.bias_v,\n",
    "                self.add_zero_attn,\n",
    "                self.dropout,\n",
    "                self.out_proj.weight,\n",
    "                self.out_proj.bias,\n",
    "                training=self.training,\n",
    "                key_padding_mask=key_padding_mask,\n",
    "                need_weights=need_weights,\n",
    "                attn_mask=attn_mask,\n",
    "            )\n",
    "\n",
    "\n",
    "def _get_clones(module, N):\n",
    "    return ModuleList([copy.deepcopy(module) for i in range(N)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run {UNIQUE_PATH}/model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tNCaSlrFDGA6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting /lustre04/scratch/strokach/workspace/proteinsolver/notebooks/sudoku_train/4264619e/stats.py\n"
     ]
    }
   ],
   "source": [
    "%%file {UNIQUE_PATH}/stats.py\n",
    "import atexit\n",
    "import csv\n",
    "import time\n",
    "import warnings\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class Stats:\n",
    "    epoch: int\n",
    "    step: int\n",
    "    batch_size: int\n",
    "    echo: bool\n",
    "    total_loss: float\n",
    "    num_correct_preds: int\n",
    "    num_preds: int\n",
    "    num_correct_preds_missing: int\n",
    "    num_preds_missing: int\n",
    "    num_correct_preds_missing_valid: int\n",
    "    num_preds_missing_valid: int\n",
    "    num_correct_preds_missing_valid_old: int\n",
    "    num_preds_missing_valid_old: int\n",
    "    start_time: float\n",
    "\n",
    "    def __init__(\n",
    "        self, *, epoch=0, step=0, batch_size=1, filename=None, echo=True, tb_writer=None\n",
    "    ):\n",
    "        self.epoch = epoch\n",
    "        self.step = step\n",
    "        self.batch_size = batch_size\n",
    "        self.echo = echo\n",
    "        self.tb_writer = tb_writer\n",
    "        self.prev = {}\n",
    "        self.init_parameters()\n",
    "\n",
    "        if filename:\n",
    "            self.filehandle = open(filename, \"wt\", newline=\"\")\n",
    "            self.writer = csv.DictWriter(\n",
    "                self.filehandle, list(self.stats.keys()), dialect=\"unix\"\n",
    "            )\n",
    "            atexit.register(self.filehandle.close)\n",
    "        else:\n",
    "            self.filehandle = None\n",
    "            self.writer = None\n",
    "\n",
    "    def init_parameters(self):\n",
    "        self.num_steps = 0\n",
    "        self.total_loss = 0\n",
    "        self.num_correct_preds = 0\n",
    "        self.num_preds = 0\n",
    "        self.num_correct_preds_missing = 0\n",
    "        self.num_preds_missing = 0\n",
    "        self.num_correct_preds_missing_valid = 0\n",
    "        self.num_preds_missing_valid = 0\n",
    "        self.num_correct_preds_missing_valid_old = 0\n",
    "        self.num_preds_missing_valid_old = 0\n",
    "        self.start_time = time.perf_counter()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.prev = self.stats\n",
    "        self.init_parameters()\n",
    "\n",
    "    @property\n",
    "    def header(self):\n",
    "        return \"\".join(to_fixed_width(self.stats.keys()))\n",
    "\n",
    "    @property\n",
    "    def row(self):\n",
    "        return \"\".join(to_fixed_width(self.stats.values(), 4))\n",
    "\n",
    "    @property\n",
    "    def stats(self):\n",
    "        with warnings.catch_warnings():\n",
    "            warnings.simplefilter(\"ignore\")\n",
    "            return {\n",
    "                \"epoch\": self.epoch,\n",
    "                \"step\": self.step,\n",
    "                \"datapoint\": self.datapoint,\n",
    "                \"avg_loss\": np.float64(1) * self.total_loss / self.num_steps,\n",
    "                \"accuracy\": np.float64(1) * self.num_correct_preds / self.num_preds,\n",
    "                \"accuracy_m\": np.float64(1)\n",
    "                * self.num_correct_preds_missing\n",
    "                / self.num_preds_missing,\n",
    "                \"accuracy_mv\": self.accuracy_mv,\n",
    "                \"accuracy_mv_old\": self.accuracy_mv_old,\n",
    "                \"time_elapsed\": time.perf_counter() - self.start_time,\n",
    "            }\n",
    "\n",
    "    @property\n",
    "    def accuracy_mv(self):\n",
    "        return (\n",
    "            np.float64(1)\n",
    "            * self.num_correct_preds_missing_valid\n",
    "            / self.num_preds_missing_valid\n",
    "        )\n",
    "\n",
    "    @property\n",
    "    def accuracy_mv_old(self):\n",
    "        return (\n",
    "            np.float64(1)\n",
    "            * self.num_correct_preds_missing_valid_old\n",
    "            / self.num_preds_missing_valid_old\n",
    "        )\n",
    "\n",
    "    @property\n",
    "    def datapoint(self):\n",
    "        return self.step * self.batch_size\n",
    "\n",
    "    def write_header(self):\n",
    "        if self.echo:\n",
    "            print(self.header)\n",
    "        if self.writer is not None:\n",
    "            self.writer.writeheader()\n",
    "\n",
    "    def write_row(self):\n",
    "        if self.echo:\n",
    "            print(self.row, end=\"\\r\")\n",
    "        if self.writer is not None:\n",
    "            self.writer.writerow(self.stats)\n",
    "            self.filehandle.flush()\n",
    "        if self.tb_writer is not None:\n",
    "            stats = self.stats\n",
    "            datapoint = stats.pop(\"datapoint\")\n",
    "            for key, value in stats.items():\n",
    "                self.tb_writer.add_scalar(key, value, datapoint)\n",
    "            self.tb_writer.flush()\n",
    "\n",
    "\n",
    "def to_fixed_width(lst, precision=None):\n",
    "    lst = [round(l, precision) if isinstance(l, float) else l for l in lst]\n",
    "    return [f\"{l: <18}\" for l in lst]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run {UNIQUE_PATH}/stats.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_stats_on_missing(x, y, output):\n",
    "    mask = (x == 9).squeeze()\n",
    "    if not mask.any():\n",
    "        return 0.0, 0.0\n",
    "    output_missing = output[mask]\n",
    "    _, predicted_missing = torch.max(output_missing.data, 1)\n",
    "    return (predicted_missing == y[mask]).sum().item(), len(predicted_missing)\n",
    "\n",
    "\n",
    "\n",
    "from contextlib import contextmanager\n",
    "@contextmanager\n",
    "def eval_net(net: nn.Module):\n",
    "    training = net.training\n",
    "    try:\n",
    "        net.train(False)\n",
    "        yield\n",
    "    finally:\n",
    "        net.train(training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(6, 200, 162)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_size = 6\n",
    "info_size = 200\n",
    "hidden_size = 162\n",
    "checkpoint_size = 100_000\n",
    "\n",
    "batch_size, info_size, hidden_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('/lustre04/scratch/strokach/workspace/proteinsolver/notebooks/sudoku_train/runs/4264619e')"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensorboard_path = NOTEBOOK_PATH.joinpath(\"runs\", UNIQUE_PATH.name)\n",
    "tensorboard_path.mkdir(exist_ok=True)\n",
    "tensorboard_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0,\n",
       " 66667,\n",
       " 400002,\n",
       " PosixPath('/lustre04/scratch/strokach/workspace/proteinsolver/notebooks/sudoku_train/4264619e/e0-s66667-d400002-amv07144.state'))"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "last_epoch = None\n",
    "last_step = None\n",
    "last_datapoint = None\n",
    "last_state_file = None\n",
    "\n",
    "if CONTINUE_PREVIOUS:\n",
    "    for path in UNIQUE_PATH.glob(\"*.state\"):\n",
    "        e, s, d, amv = path.name.split(\"-\")\n",
    "        datapoint = int(d.strip(\"d\"))\n",
    "        if last_datapoint is None or datapoint >= last_datapoint:\n",
    "            last_datapoint = datapoint\n",
    "            last_epoch = int(e.strip(\"e\"))\n",
    "            last_step = int(s.strip(\"s\"))\n",
    "            last_state_file = path\n",
    "        \n",
    "last_epoch, last_step, last_datapoint, last_state_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 6380
    },
    "colab_type": "code",
    "id": "azxPKTSKwcE8",
    "outputId": "09ed9b8b-25b9-49ff-8d02-83a4ec0ecba1"
   },
   "outputs": [],
   "source": [
    "net = Net(\n",
    "    x_input_size=13,\n",
    "    adj_input_size=3,\n",
    "    hidden_size=hidden_size,\n",
    "    output_size=9,\n",
    "    batch_size=batch_size,\n",
    ").to(device)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(net.parameters(), lr=0.0001)\n",
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, \"max\", verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 6380
    },
    "colab_type": "code",
    "id": "azxPKTSKwcE8",
    "outputId": "09ed9b8b-25b9-49ff-8d02-83a4ec0ecba1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded network state file.\n"
     ]
    }
   ],
   "source": [
    "if CONTINUE_PREVIOUS:\n",
    "    net.load_state_dict(torch.load(last_state_file))\n",
    "    print(\"Loaded network state file.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 6380
    },
    "colab_type": "code",
    "id": "azxPKTSKwcE8",
    "outputId": "09ed9b8b-25b9-49ff-8d02-83a4ec0ecba1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch             step              datapoint         avg_loss          accuracy          accuracy_m        accuracy_mv       accuracy_mv_old   time_elapsed      \n"
     ]
    }
   ],
   "source": [
    "stats = Stats(\n",
    "    epoch=last_epoch if CONTINUE_PREVIOUS else 0,\n",
    "    step=last_step if CONTINUE_PREVIOUS else 0,\n",
    "    batch_size=batch_size,\n",
    "    filename=UNIQUE_PATH.joinpath(\"training.log\"),\n",
    "    echo=True,\n",
    "    tb_writer=torch.utils.tensorboard.writer.SummaryWriter(\n",
    "        log_dir=tensorboard_path.with_suffix(f\".xxx\"),\n",
    "        purge_step=(last_datapoint if CONTINUE_PREVIOUS else None),\n",
    "    ),\n",
    ")\n",
    "stats.write_header()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 6380
    },
    "colab_type": "code",
    "id": "azxPKTSKwcE8",
    "outputId": "09ed9b8b-25b9-49ff-8d02-83a4ec0ecba1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0                 72534             435204            0.4503            0.8049            0.7045            0.7195            1.0               35.2388           \r"
     ]
    }
   ],
   "source": [
    "datasets[f\"sudoku_valid_0\"].reset()\n",
    "valid_0_data = list(itertools.islice(datasets[f\"sudoku_valid_0\"], 300))\n",
    "valid_old_data = list(itertools.islice(datasets[f\"sudoku_valid_old\"], 300))\n",
    "tmp_data = valid_0_data[0].to(device)\n",
    "edge_index = tmp_data.edge_index\n",
    "edge_attr = tmp_data.edge_attr\n",
    "\n",
    "net = net.train()\n",
    "for epoch in range(stats.epoch, 100_000):\n",
    "    stats.epoch = epoch\n",
    "    train_dataloader = DataLoader(\n",
    "        datasets[f\"sudoku_train_{epoch}\"],\n",
    "        shuffle=False,\n",
    "        num_workers=1,\n",
    "        batch_size=batch_size,\n",
    "        drop_last=True,\n",
    "    )\n",
    "    for data in train_dataloader:\n",
    "        stats.step += 1\n",
    "        if CONTINUE_PREVIOUS and stats.step <= last_step:\n",
    "            continue\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        data = data.to(device)\n",
    "        output = net(data.x, data.edge_index, data.edge_attr)\n",
    "\n",
    "        loss = criterion(output, data.y)\n",
    "        loss.backward()\n",
    "\n",
    "        stats.total_loss += loss.detach().item()\n",
    "        stats.num_steps += 1\n",
    "\n",
    "        # Accuracy for all\n",
    "        _, predicted = torch.max(output.data, 1)\n",
    "        stats.num_correct_preds += (predicted == data.y).sum().item()\n",
    "        stats.num_preds += len(predicted)\n",
    "\n",
    "        # Accuracy for missing only\n",
    "        num_correct, num_total = get_stats_on_missing(data.x, data.y, output)\n",
    "        stats.num_correct_preds_missing += num_correct\n",
    "        stats.num_preds_missing += num_total\n",
    "\n",
    "        optimizer.step()\n",
    "\n",
    "        if (stats.datapoint % info_size) < batch_size:\n",
    "            for j, data in enumerate(valid_0_data):\n",
    "                data = data.to(device)\n",
    "\n",
    "                with torch.no_grad() and eval_net(net):\n",
    "                    output = net(data.x, data.edge_index, data.edge_attr)\n",
    "\n",
    "                num_correct, num_total = get_stats_on_missing(data.x, data.y, output)\n",
    "                stats.num_correct_preds_missing_valid += num_correct\n",
    "                stats.num_preds_missing_valid += num_total\n",
    "\n",
    "            for j, data in enumerate(valid_old_data):\n",
    "                data = data.to(device)\n",
    "\n",
    "                with torch.no_grad() and eval_net(net):\n",
    "                    output = net(data.x, data.edge_index, edge_attr)\n",
    "\n",
    "                num_correct, num_total = get_stats_on_missing(data.x, data.y, output)\n",
    "                stats.num_correct_preds_missing_valid_old += num_correct\n",
    "                stats.num_preds_missing_valid_old += num_total\n",
    "\n",
    "            stats.write_row()\n",
    "            stats.reset_parameters()\n",
    "\n",
    "        if (stats.datapoint % checkpoint_size) < batch_size:\n",
    "            output_filename = (\n",
    "                f\"e{stats.epoch}-s{stats.step}-d{stats.datapoint}\"\n",
    "                f\"-amv{str(round(stats.prev['accuracy_mv'], 4)).replace('.', '')}.state\"\n",
    "            )\n",
    "            torch.save(net.state_dict(), UNIQUE_PATH.joinpath(output_filename))\n",
    "\n",
    "    scheduler.step(stats.prev[\"accuracy_mv\"])\n",
    "    output_filename = (\n",
    "        f\"e{stats.epoch}-s{stats.step}-d{stats.datapoint}\"\n",
    "        f\"-amv{str(round(stats.prev['accuracy_mv'], 4)).replace('.', '')}.state\"\n",
    "    )\n",
    "    torch.save(net.state_dict(), UNIQUE_PATH.joinpath(output_filename))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "1fXJz1l9HvhW",
    "u9q4iTifHvhl"
   ],
   "name": "Copy of 2019-03-30-sudoku-4xEdgeConv-09862+.ipynb",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}