notebooks/04_protein_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "- 5 layer graph-conv + attn.\n",
    "- batch size = 1.\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "t_gnhLD0HvhU",
    "toc-hr-collapsed": true,
    "toc-nb-collapsed": true
   },
   "source": [
    "## Install dependencies (Google Colab only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import google.colab\n",
    "    GOOGLE_COLAB = True\n",
    "except ImportError:\n",
    "    GOOGLE_COLAB = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 220
    },
    "colab_type": "code",
    "id": "HCWw1IgNHvhm",
    "outputId": "fcb1a352-1ca8-4708-c3af-11bb0f9a3c57"
   },
   "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,
    "toc-nb-collapsed": true
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TbKxMUZWHvhq"
   },
   "outputs": [],
   "source": [
    "import atexit\n",
    "import csv\n",
    "import tempfile\n",
    "import time\n",
    "import uuid\n",
    "import warnings\n",
    "from collections import deque\n",
    "from contextlib import contextmanager\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import pyarrow\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.tensorboard\n",
    "from torch import optim\n",
    "from torch_geometric.data import DataLoader\n",
    "from torch_geometric.nn import ChebConv, EdgeConv, GATConv, GCNConv\n",
    "from torch_geometric.nn.inits import reset\n",
    "from torch_geometric.utils import add_self_loops, remove_self_loops, scatter_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 334
    },
    "colab_type": "code",
    "id": "rczy7pPiHvhs",
    "outputId": "8e4673d7-4cc7-4f8e-ffc4-de94302e5fab"
   },
   "outputs": [],
   "source": [
    "import proteinsolver\n",
    "import proteinsolver.datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4fFnAyUOHvhv"
   },
   "outputs": [],
   "source": [
    "# %load_ext autoreload\n",
    "# %autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_ROOT = Path(tempfile.gettempdir())\n",
    "DATA_ROOT = Path(\"/home/strokach/ml_data\")\n",
    "DATA_ROOT.mkdir(exist_ok=True)\n",
    "DATA_ROOT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    NOTEBOOK_PATH\n",
    "    UNIQUE_PATH\n",
    "except NameError:\n",
    "    NOTEBOOK_PATH = Path(\"protein_train\").resolve()\n",
    "    NOTEBOOK_PATH.mkdir(exist_ok=True)\n",
    "    unique_id = uuid.uuid4().hex[:8]\n",
    "    UNIQUE_PATH = NOTEBOOK_PATH.joinpath(unique_id)\n",
    "    UNIQUE_PATH.mkdir()\n",
    "NOTEBOOK_PATH, UNIQUE_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATAPKG_OUTPUT_DIR = Path(f\"~/datapkg_output_dir\").expanduser().resolve()\n",
    "DATAPKG_OUTPUT_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proteinsolver.settings.data_url = DATAPKG_OUTPUT_DIR.as_posix()\n",
    "proteinsolver.settings.data_url"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dsEY3dtLHvhy",
    "toc-hr-collapsed": true,
    "toc-nb-collapsed": true
   },
   "source": [
    "# Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(10):\n",
    "    dataset_name = f\"protein_train_{i}\"\n",
    "    datasets[dataset_name] = proteinsolver.datasets.ProteinDataset2(\n",
    "        root=DATA_ROOT.joinpath(dataset_name), subset=f\"protein_train_{i}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets[\"protein_valid\"] = proteinsolver.datasets.ProteinInMemoryDataset(root=DATA_ROOT / \"protein_valid\", subset=\"valid\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets[\"protein_test\"] = proteinsolver.datasets.ProteinInMemoryDataset(root=DATA_ROOT / \"protein_test\", subset=\"test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file = \"/home/kimlab1/database_data/datapkg_output_dir/adjacency-net-v2/master/validation_dataset_wdistances/adjacency_matrix.parquet/database_id=G3DSA%3A2.40.155.10/part-00000-d5e89475-69dd-45c6-9a80-5bf751fce422-c000.snappy.parquet\"\n",
    "datasets[\"protein_gfp\"] = proteinsolver.datasets.ProteinInMemoryDataset(\n",
    "    root=DATA_ROOT / \"protein_gfp\", subset=\"gfp\", data_url=file\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e3kWuRaxr89h"
   },
   "source": [
    "# Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%file {UNIQUE_PATH}/model.py\n",
    "import atexit\n",
    "import copy\n",
    "import csv\n",
    "import tempfile\n",
    "import time\n",
    "import warnings\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 add_self_loops, remove_self_loops, scatter_, to_dense_adj, to_dense_batch\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.nn = nn.Sequential(\n",
    "            #\n",
    "            nn.Linear(hidden_size * 3, 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",
    "        # 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",
    "        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, num_heads):\n",
    "        super().__init__()\n",
    "        self.attn = nn.MultiheadAttention(hidden_size, num_heads)\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 = 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__(self, x_input_size, adj_input_size, hidden_size, output_size):\n",
    "        super().__init__()\n",
    "\n",
    "        self.embed_x = nn.Sequential(\n",
    "            nn.Embedding(x_input_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, hidden_size),\n",
    "            nn.LayerNorm(hidden_size),\n",
    "            #             nn.ReLU(),\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 = 5\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, 8)\n",
    "        self.attns = _get_clones(attn, N)\n",
    "\n",
    "        self.linear_out = nn.Linear(hidden_size, output_size)\n",
    "        \n",
    "        self.register_buffer('batch', torch.zeros(10000, dtype=torch.int64))\n",
    "\n",
    "    def forward(self, x, edge_index, edge_attr):\n",
    "\n",
    "        x = self.embed_x(x)\n",
    "#         edge_index, _ = add_self_loops(edge_index)  # We should remove self loops in this case!\n",
    "        edge_attr = self.embed_adj(edge_attr) if edge_attr is not None else None\n",
    "\n",
    "        i = 0\n",
    "        edge_attr_out = self.edge_convs[i](x, edge_index, edge_attr)\n",
    "        edge_attr = edge_attr + F.dropout(edge_attr_out, 0.1)\n",
    "        edge_attr = self.adj_norms_1[i](edge_attr)\n",
    "\n",
    "        x_out = self.attns[i](x, edge_index, self.adj_norms_0[i](edge_attr_out), self.batch[:x.size(0)])\n",
    "        x = x + F.dropout(x_out, 0.1)\n",
    "        x = self.x_norms_1[i](x)\n",
    "    \n",
    "\n",
    "        i = 1\n",
    "        edge_attr_out = self.edge_convs[i](x, edge_index, edge_attr)\n",
    "        edge_attr = edge_attr + F.dropout(edge_attr_out, 0.1)\n",
    "        edge_attr = self.adj_norms_1[i](edge_attr)\n",
    "\n",
    "        x_out = self.attns[i](x, edge_index, self.adj_norms_0[i](edge_attr_out), self.batch[:x.size(0)])\n",
    "        x = x + F.dropout(x_out, 0.1)\n",
    "        x = self.x_norms_1[i](x)\n",
    "    \n",
    "\n",
    "        i = 2\n",
    "        edge_attr_out = self.edge_convs[i](x, edge_index, edge_attr)\n",
    "        edge_attr = edge_attr + F.dropout(edge_attr_out, 0.1)\n",
    "        edge_attr = self.adj_norms_1[i](edge_attr)\n",
    "\n",
    "        x_out = self.attns[i](x, edge_index, self.adj_norms_0[i](edge_attr_out), self.batch[:x.size(0)])\n",
    "        x = x + F.dropout(x_out, 0.1)\n",
    "        x = self.x_norms_1[i](x)\n",
    "    \n",
    "\n",
    "        i = 3\n",
    "        edge_attr_out = self.edge_convs[i](x, edge_index, edge_attr)\n",
    "        edge_attr = edge_attr + F.dropout(edge_attr_out, 0.1)\n",
    "        edge_attr = self.adj_norms_1[i](edge_attr)\n",
    "\n",
    "        x_out = self.attns[i](x, edge_index, self.adj_norms_0[i](edge_attr_out), self.batch[:x.size(0)])\n",
    "        x = x + F.dropout(x_out, 0.1)\n",
    "        x = self.x_norms_1[i](x)\n",
    "\n",
    "\n",
    "        i = 4\n",
    "        edge_attr_out = self.edge_convs[i](x, edge_index, edge_attr)\n",
    "        edge_attr = edge_attr + F.dropout(edge_attr_out, 0.1)\n",
    "        edge_attr = self.adj_norms_1[i](edge_attr)\n",
    "\n",
    "        x_out = self.attns[i](x, edge_index, self.adj_norms_0[i](edge_attr_out), self.batch[:x.size(0)])\n",
    "        x = x + F.dropout(x_out, 0.1)\n",
    "        x = self.x_norms_1[i](x)\n",
    "\n",
    "\n",
    "        x = self.linear_out(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "def _get_clones(module, N):\n",
    "    return ModuleList([copy.deepcopy(module) for i in range(N)])\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]\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",
    "    start_time: float\n",
    "\n",
    "    def __init__(self, *, epoch=0, step=0, batch_size=1, filename=None, echo=True, tb_writer=None):\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.reset_parameters()\n",
    "\n",
    "        if filename:\n",
    "            self.filehandle = open(filename, \"wt\", newline=\"\")\n",
    "            self.writer = csv.DictWriter(self.filehandle, list(self.stats.keys()), dialect=\"unix\")\n",
    "            self.writer.writeheader()\n",
    "            atexit.register(self.filehandle.close)\n",
    "        else:\n",
    "            self.filehandle = None\n",
    "            self.writer = None\n",
    "\n",
    "    def reset_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.start_time = time.perf_counter()\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) * self.num_correct_preds_missing / self.num_preds_missing,\n",
    "                \"accuracy_mv\": self.accuracy_mv,\n",
    "                \"time_elapsed\": time.perf_counter() - self.start_time,\n",
    "            }\n",
    "\n",
    "    @property\n",
    "    def accuracy_mv(self):\n",
    "        return np.float64(1) * self.num_correct_preds_missing_valid / self.num_preds_missing_valid\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)\n",
    "        if self.writer is not None:\n",
    "            self.writer.writerow(self.stats)\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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run {UNIQUE_PATH}/model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters\n",
    "batch_size = 4\n",
    "num_features = 20\n",
    "adj_input_size = 2\n",
    "hidden_size = 128\n",
    "frac_present = 0.5\n",
    "frac_present_valid = frac_present\n",
    "info_size= 1024\n",
    "\n",
    "# dataloaders = {\n",
    "#     \"train_0\": DataLoader(protein_dataset_train_0, shuffle=True, num_workers=4, batch_size=batch_size, drop_last=True),\n",
    "#     \"train_1\": DataLoader(protein_dataset_train_1, shuffle=True, num_workers=4, batch_size=batch_size, drop_last=True),\n",
    "#     \"train_2\": None,\n",
    "#     \"valid\": DataLoader(protein_dataset_valid[:128], shuffle=False, num_workers=4, batch_size=1, drop_last=False),\n",
    "#     \"test\": DataLoader(protein_dataset_test[:128], shuffle=False, num_workers=4, batch_size=1, drop_last=False),\n",
    "\n",
    "#     \"gfp_train\": DataLoader(protein_dataset_test[:928], shuffle=True, num_workers=4, batch_size=batch_size, drop_last=True),\n",
    "#     \"gfp_valid\": DataLoader(protein_dataset_test[928:], shuffle=False, num_workers=4, batch_size=1, drop_last=False),\n",
    "# }"
   ]
  },
  {
   "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": [],
   "source": [
    "@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)\n",
    "\n",
    "\n",
    "def get_stats_on_missing(data, output):\n",
    "    mask = (data.x == num_features).squeeze()\n",
    "    output_missing = output[mask]\n",
    "    _, predicted_missing = torch.max(output_missing.data, 1)\n",
    "    return (predicted_missing == data.y[mask]).sum().item(), len(predicted_missing)\n",
    "\n",
    "\n",
    "def get_data_x(data, frac_present):\n",
    "    x = torch.where(\n",
    "        torch.rand(data.y.size(0), device=data.y.device) < frac_present,\n",
    "        data.y,\n",
    "        torch.ones(1, dtype=torch.long, device=data.y.device) * num_features,\n",
    "    )\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorboard_path = NOTEBOOK_PATH.joinpath(\"runs\", UNIQUE_PATH.name)"
   ]
  },
  {
   "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": [],
   "source": [
    "continue_previous = False\n",
    "\n",
    "if not continue_previous:\n",
    "    net = Net(\n",
    "        x_input_size=num_features + 1, adj_input_size=adj_input_size, hidden_size=hidden_size, output_size=num_features\n",
    "    )\n",
    "    net = net.to(device)\n",
    "    stats = Stats(\n",
    "        epoch=0,\n",
    "        step=0,\n",
    "        batch_size=batch_size,\n",
    "        filename=UNIQUE_PATH.joinpath(\"training.log\"),\n",
    "        echo=True,\n",
    "        tb_writer=torch.utils.tensorboard.writer.SummaryWriter(log_dir=tensorboard_path.with_suffix(f\".xxx\")),\n",
    "    )\n",
    "    optimizer = optim.Adam(net.parameters(), lr=0.0001)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, \"max\", verbose=True)\n",
    "\n",
    "net = net.train()\n",
    "stats.write_header()\n",
    "for epoch in range(stats.epoch + 1 if continue_previous else 0, 100_000):\n",
    "    stats.epoch = epoch\n",
    "    train_dl = iter(\n",
    "        DataLoader(\n",
    "            datasets[f\"protein_train_{epoch % 10}\"], shuffle=False, num_workers=1, batch_size=batch_size, drop_last=True\n",
    "        )\n",
    "    )\n",
    "    while True:\n",
    "        try:\n",
    "            data = next(train_dl)\n",
    "        except StopIteration:\n",
    "            break\n",
    "\n",
    "        stats.step += 1\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        max_size = 700\n",
    "        if data.x.size(0) > max_size:\n",
    "            data.x = data.x[:max_size]\n",
    "            mask = (data.edge_index < max_size).all(dim=0)\n",
    "            data.edge_index = data.edge_index[:, mask]\n",
    "            data.edge_attr = data.edge_attr[mask, :]\n",
    "\n",
    "        data = data.to(device)\n",
    "        data.y = data.x\n",
    "        data.x = get_data_x(data, frac_present)\n",
    "        output = net(data.x, data.edge_index, data.edge_attr)\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, 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(\n",
    "                DataLoader(\n",
    "                    datasets[\"protein_valid\"][:128], shuffle=False, num_workers=1, batch_size=1, drop_last=True\n",
    "                )\n",
    "            ):\n",
    "                data = data.to(device)\n",
    "                data.y = data.x\n",
    "                data.x = get_data_x(data, frac_present_valid)\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, output)\n",
    "                stats.num_correct_preds_missing_valid += num_correct\n",
    "                stats.num_preds_missing_valid += num_total\n",
    "\n",
    "            #             scheduler.step(stats.stats['accuracy'])\n",
    "            stats.write_row()\n",
    "            stats.reset_parameters()\n",
    "    output_filename = f\"e{stats.epoch}-s{stats.step}-d{stats.datapoint}.state\"\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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}