CRaWl_(hyperparameter_in_progress).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CRaWl (hyperparameter in progress).ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyOtZ4rxrEqGla6loFiljhtO",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/2019mohamed/Enzyme-classification/blob/master/CRaWl_(hyperparameter_in_progress).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IcboYcH_0ly5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c7a6b067-2536-49ba-fef7-bb66c5985f78"
      },
      "source": [
        "pip install deepsnap"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting deepsnap\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/cb/a9/d785dd32167f7d2bf9d40292098c4f3f3783c3e83ade1a4e481c46ef0a80/deepsnap-0.2.0-py3-none-any.whl (76kB)\n",
            "\u001b[K     |████████████████████████████████| 81kB 3.3MB/s \n",
            "\u001b[?25hRequirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (from deepsnap) (1.8.1+cu101)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from deepsnap) (1.19.5)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.7/dist-packages (from deepsnap) (2.5.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch->deepsnap) (3.7.4.3)\n",
            "Requirement already satisfied: decorator<5,>=4.3 in /usr/local/lib/python3.7/dist-packages (from networkx->deepsnap) (4.4.2)\n",
            "Installing collected packages: deepsnap\n",
            "Successfully installed deepsnap-0.2.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x-fJHq--3KZf",
        "outputId": "42b1cd1c-0896-4295-9f8f-06147c3fc165"
      },
      "source": [
        "#! pip install torch-scatter \n",
        "! pip install torch-sparse \n",
        "#! pip install torch-cluster \n",
        "#! pip install torch-spline-conv \n",
        "#! pip install torch-geometric"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting torch-sparse\n",
            "  Downloading https://files.pythonhosted.org/packages/9a/86/699eb78ea7ce259da7f9953858ec2a064e4e5f5e6bf7de53aa6c8bb8b9a8/torch_sparse-0.6.9.tar.gz\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from torch-sparse) (1.4.1)\n",
            "Requirement already satisfied: numpy>=1.13.3 in /usr/local/lib/python3.7/dist-packages (from scipy->torch-sparse) (1.19.5)\n",
            "Building wheels for collected packages: torch-sparse\n",
            "  Building wheel for torch-sparse (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for torch-sparse: filename=torch_sparse-0.6.9-cp37-cp37m-linux_x86_64.whl size=432311 sha256=0c2448cea1c13a1f57de39d23b713df1ad2e16c07d74cbcfbfe8cd91b2f0f29c\n",
            "  Stored in directory: /root/.cache/pip/wheels/6c/9f/3e/8813c1f7e87d12b01b830d3506e7f3f6b044d9d53769057ddb\n",
            "Successfully built torch-sparse\n",
            "Installing collected packages: torch-sparse\n",
            "Successfully installed torch-sparse-0.6.9\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XLUPLRLO7fCl",
        "outputId": "20dda116-c71f-4cd9-c503-1dc39d1ea8e0"
      },
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import torch\n",
        "from torch_scatter import scatter_mean, scatter_sum\n",
        "\n",
        "MAXINT = np.iinfo(np.int64).max\n",
        "#https://github.com/toenshoff/CRaWl\n",
        "import torch\n",
        "import numpy as np\n",
        "import torch_geometric as pygeo\n",
        "from torch.nn import functional as F\n",
        "from torch.utils.data import DataLoader\n",
        "from torch_geometric.data import Data\n",
        "from torch_scatter import scatter_sum\n",
        "\n",
        "\n",
        "def preproc(data):\n",
        "    \"\"\" Preprocess Pytorch Geometric data objects to be used with our walk generator \"\"\"\n",
        "\n",
        "    if data.num_edges == 0:\n",
        "        return data\n",
        "\n",
        "    if not data.is_coalesced():\n",
        "        data.coalesce()\n",
        "\n",
        "    if data.num_node_features == 0:\n",
        "        data.x = torch.zeros((data.num_nodes, 1), dtype=torch.float32)\n",
        "\n",
        "    if data.num_edge_features == 0:\n",
        "        data.edge_attr = torch.zeros((data.num_edges, 1), dtype=torch.float32)\n",
        "\n",
        "    edge_idx = data.edge_index\n",
        "    edge_feat = data.edge_attr\n",
        "    node_feat = data.x\n",
        "\n",
        "    # remove isolated nodes\n",
        "    if data.contains_isolated_nodes():\n",
        "        edge_idx, edge_feat, mask = pygeo.utils.remove_isolated_nodes(edge_idx, edge_feat, data.num_nodes)\n",
        "        node_feat = node_feat[mask]\n",
        "\n",
        "    # Enforce undirected graphs\n",
        "    if not pygeo.utils.is_undirected(edge_idx):\n",
        "        x = edge_feat.detach().numpy()\n",
        "        e = edge_idx.detach().numpy()\n",
        "        x_map = {(e[0,i], e[1,i]): x[i] for i in range(e.shape[1])}\n",
        "        edge_idx = pygeo.utils.to_undirected(edge_idx)\n",
        "        e = edge_idx.detach().numpy()\n",
        "        x = [x_map[(e[0,i], e[1,i])] if (e[0,i], e[1,i]) in x_map.keys() else x_map[(e[1,i], e[0,i])] for i in range(e.shape[1])]\n",
        "        edge_feat = torch.tensor(x)\n",
        "\n",
        "    data.edge_index = edge_idx\n",
        "    data.edge_attr = edge_feat\n",
        "    data.x = node_feat\n",
        "\n",
        "    if not torch.is_tensor(data.y):\n",
        "        data.y = torch.tensor(data.y)\n",
        "    data.y = data.y.view(1, -1)\n",
        "\n",
        "    order = node_feat.shape[0]\n",
        "\n",
        "    # create bitwise encoding of adjacency matrix using 64-bit integers\n",
        "    data.node_id = torch.arange(0, order)\n",
        "    bit_id = torch.zeros((order, order // 63 + 1), dtype=torch.int64)\n",
        "    bit_id[data.node_id, data.node_id // 63] = torch.tensor(1) << data.node_id % 63\n",
        "    data.adj_bits = scatter_sum(bit_id[edge_idx[0]], edge_idx[1], dim=0)\n",
        "\n",
        "    # compute node offsets in the adjacency list\n",
        "    data.degrees = pygeo.utils.degree(edge_idx[0], dtype=torch.int64)\n",
        "    adj_offset = torch.zeros((order,), dtype=torch.int64)\n",
        "    adj_offset[1:] = torch.cumsum(data.degrees, dim=0)[:-1]\n",
        "    data.adj_offset = adj_offset\n",
        "\n",
        "    return data\n",
        "\n",
        "\n",
        "def merge_batch(graph_data):\n",
        "    \"\"\" Custom function to collate preprocessed data objects in the data loader \"\"\"\n",
        "\n",
        "    adj_offset = [d.adj_offset for d in graph_data]\n",
        "    degrees = [d.degrees for d in graph_data]\n",
        "    edge_idx = [d.edge_index for d in graph_data]\n",
        "\n",
        "    num_nodes = torch.tensor([d.shape[0] for d in degrees])\n",
        "    num_edges = torch.tensor([e.shape[1] for e in edge_idx])\n",
        "\n",
        "    x_node = torch.cat([d.x for d in graph_data], dim=0)\n",
        "    x_edge = torch.cat([d.edge_attr for d in graph_data], dim=0)\n",
        "    x_edge = x_edge.view(x_edge.shape[0], -1)\n",
        "\n",
        "    adj_offset = torch.cat(adj_offset)\n",
        "    degrees = torch.cat(degrees)\n",
        "    edge_idx = torch.cat(edge_idx, dim=1)\n",
        "\n",
        "    node_graph_idx = torch.cat([i * torch.ones(x, dtype=torch.int64) for i, x in enumerate(num_nodes)])\n",
        "    edge_graph_idx = torch.cat([i * torch.ones(x, dtype=torch.int64) for i, x in enumerate(num_edges)])\n",
        "\n",
        "    node_shift = torch.zeros((len(graph_data),), dtype=torch.int64)\n",
        "    edge_shift = torch.zeros((len(graph_data),), dtype=torch.int64)\n",
        "    node_shift[1:] = torch.cumsum(num_nodes, dim=0)[:-1]\n",
        "    edge_shift[1:] = torch.cumsum(num_edges, dim=0)[:-1]\n",
        "\n",
        "    adj_offset += edge_shift[node_graph_idx]\n",
        "    edge_idx += node_shift[edge_graph_idx].view(1, -1)\n",
        "\n",
        "    graph_offset = node_shift\n",
        "\n",
        "    adj_bits = [d.adj_bits for d in graph_data]\n",
        "    max_enc_length = np.max([p.shape[1] for p in adj_bits])\n",
        "    adj_bits = torch.cat([F.pad(b, (0,max_enc_length-b.shape[1],0,0), 'constant', 0) for b in adj_bits], dim=0)\n",
        "\n",
        "    node_id = torch.cat([d.node_id for d in graph_data], dim=0)\n",
        "\n",
        "    y = torch.cat([d.y for d in graph_data], dim=0)\n",
        "\n",
        "    data = Data(x=x_node, edge_index=edge_idx, edge_attr=x_edge, y=y)\n",
        "    data.batch = node_graph_idx\n",
        "    data.adj_offset = adj_offset\n",
        "    data.degrees = degrees\n",
        "    data.graph_offset = graph_offset\n",
        "    data.order = num_nodes\n",
        "    data.node_id = node_id\n",
        "    data.adj_bits = adj_bits\n",
        "    return data\n",
        "\n",
        "\n",
        "class CRaWlLoader(DataLoader):\n",
        "    \"\"\" Custom Loader for our data objects \"\"\"\n",
        "\n",
        "    def __init__(self, dataset, batch_size=1, shuffle=False, **kwargs):\n",
        "        super(CRaWlLoader, self).__init__(dataset, batch_size, shuffle, collate_fn=merge_batch, **kwargs)\n",
        "\n",
        "class Walker(torch.nn.Module):\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super(Walker, self).__init__()\n",
        "\n",
        "        self.steps = config['steps']\n",
        "        #self.train_start_ratio = config['train_start_ratio']\n",
        "        self.win_size = config['win_size']\n",
        "\n",
        "        self.compute_id = 'compute_id_feat' not in config.keys() or config['compute_id_feat']\n",
        "        self.compute_adj = 'compute_adj_feat' not in config.keys() or config['compute_adj_feat']\n",
        "\n",
        "        self.struc_feat_dim = 0\n",
        "        if self.compute_id:\n",
        "            self.struc_feat_dim += self.win_size\n",
        "        if self.compute_adj:\n",
        "            self.struc_feat_dim += self.win_size - 1\n",
        "\n",
        "        self.choice_fct = Walker.uniform_choice if config['walk_mode'] == 'uniform' else Walker.nb_choice\n",
        "\n",
        "    @staticmethod\n",
        "    def sample_start(start_p, graph_idx, graph_offset, order, device):\n",
        "        \"\"\"\n",
        "        Randomly sample start nodes\n",
        "        :param start_p: Probability of starting a walk at a node\n",
        "        :param graph_idx: Assignment of nodes to graphs\n",
        "        :param graph_offset: Node list offset of each graph in the batch\n",
        "        :param order: Nuber of nodes in each graph of the batch\n",
        "        :param device: device to construct tensors on\n",
        "        :return: A tensor of start vertices (index list) and an assignment to the graphs in the batch\n",
        "        \"\"\"\n",
        "\n",
        "        num_graphs = order.shape[0]\n",
        "        num_nodes = graph_idx.shape[0]\n",
        "        num_walks = int(np.ceil(start_p * num_nodes))\n",
        "        num_extra = num_walks - num_graphs\n",
        "\n",
        "        idx = graph_offset + (torch.randint(0, MAXINT, (num_graphs,), device=device) % order)\n",
        "        idx = torch.cat([idx, torch.randperm(num_nodes, device=device)[:num_extra]])\n",
        "\n",
        "        choices = torch.randint(0, MAXINT, (num_walks,), device=device)\n",
        "        start_graph = graph_idx[idx]\n",
        "        start = graph_offset[start_graph] + (choices % order[start_graph])\n",
        "\n",
        "        del idx, choices\n",
        "        return start, start_graph\n",
        "\n",
        "    @staticmethod\n",
        "    def uniform_choice(i, walks, adj_nodes, adj_offset, choices, degrees, nb_degrees):\n",
        "        \"\"\"\n",
        "        :param i: Index of the current step\n",
        "        :param walks: Tensor of vertices in the walk\n",
        "        :param adj_nodes: Adjacency List\n",
        "        :param adj_offset: Node offset in the adjacency list\n",
        "        :param choices: Cache of random integers\n",
        "        :param degrees: Degree of each node\n",
        "        :param nb_degrees: Reduced degrees for no-backtrack walks (not used here)\n",
        "        :return: A list of a chosen outgoing edge for each walk\n",
        "        \"\"\"\n",
        "        return adj_offset[walks[i]] + (choices[i] % degrees[walks[i]])\n",
        "\n",
        "    @staticmethod\n",
        "    def nb_choice(i, walks, adj_nodes, adj_offset, choices, degrees, nb_degrees):\n",
        "        \"\"\"\n",
        "        :param i: Index of the current step\n",
        "        :param walks: Tensor of vertices in the walk\n",
        "        :param adj_nodes: Adjacency List\n",
        "        :param adj_offset: Node offset in the adjacency list\n",
        "        :param choices: Cache of random integers\n",
        "        :param degrees: Degree of each node\n",
        "        :param nb_degrees: Reduced degrees for no-backtrack walks\n",
        "        :return: A list of a chosen outgoing edge for each walk\n",
        "        \"\"\"\n",
        "\n",
        "        # do uniform step\n",
        "        old_nodes = walks[i-1]\n",
        "        cur_nodes = walks[i]\n",
        "        edge_idx = choices[i] % degrees[cur_nodes]\n",
        "        chosen_edges = adj_offset[cur_nodes] + edge_idx\n",
        "        new_nodes = adj_nodes[chosen_edges]\n",
        "\n",
        "        # correct backtracking\n",
        "        bt = new_nodes == old_nodes\n",
        "        if bt.max():\n",
        "            bt_nodes = walks[i][bt]\n",
        "            chosen_edges[bt] = adj_offset[bt_nodes] + (edge_idx[bt] + 1 + (choices[i][bt] % nb_degrees[bt_nodes])) % degrees[bt_nodes]\n",
        "        return chosen_edges\n",
        "\n",
        "    def sample_walks(self, data, x_edge, steps=None, start_p=1.0):\n",
        "        \"\"\"\n",
        "        :param data: Preprocessed PyTorch Geometric data object.\n",
        "        :param x_edge: Edge features\n",
        "        :param steps: Number of walk steps (if None, default from config is used)\n",
        "        :param start_p: Probability of starting a walk at each node\n",
        "        :return: The data object with the walk added as an attribute\n",
        "        \"\"\"\n",
        "\n",
        "        device = data.x.device\n",
        "\n",
        "        # get adjacency data\n",
        "        adj_nodes = data.edge_index[1]\n",
        "        adj_offset = data.adj_offset\n",
        "        degrees = data.degrees\n",
        "        node_id = data.node_id\n",
        "        adj_bits = data.adj_bits\n",
        "        graph_idx = data.batch\n",
        "        graph_offset = data.graph_offset\n",
        "        order = data.order\n",
        "\n",
        "        # use default number of steps if not specified\n",
        "        if steps is None:\n",
        "            steps = self.steps\n",
        "\n",
        "        # set dimensions\n",
        "        s = self.win_size\n",
        "        n = degrees.shape[0]\n",
        "        l = steps + 1\n",
        "\n",
        "        # sample starting nodes\n",
        "        if start_p < 1.0:\n",
        "            start, walk_graph_idx = Walker.sample_start(start_p, graph_idx, graph_offset, order, device)\n",
        "        else:\n",
        "            start = torch.arange(0, n, dtype=torch.int64).view(-1)\n",
        "\n",
        "        # init tensor to hold walk indices\n",
        "        w = start.shape[0]\n",
        "        walks = torch.zeros((l, w), dtype=torch.int64, device=device)\n",
        "        walks[0] = start\n",
        "\n",
        "        # get all random decisions at once (faster then individual calls)\n",
        "        choices = torch.randint(0, MAXINT, (steps, w), device=device)\n",
        "\n",
        "        # init feature tensors\n",
        "        edge_feat = torch.zeros((l, w, x_edge.shape[1]), dtype=torch.float32, device=device)\n",
        "\n",
        "        if self.compute_id:\n",
        "            id_enc = torch.zeros((l, s, w), dtype=torch.bool, device=device)\n",
        "\n",
        "        if self.compute_adj:\n",
        "            edges = torch.zeros((l, s, w), dtype=torch.bool, device=device)\n",
        "\n",
        "        # remove one choice of each node with deg > 1 for no_backtrack walks\n",
        "        nb_degree_mask = (degrees == 1)\n",
        "        nb_degrees = nb_degree_mask * degrees + (~nb_degree_mask) * (degrees - 1)\n",
        "\n",
        "        for i in range(steps):\n",
        "            # get next edges\n",
        "            if i == 0:\n",
        "                chosen_edges = self.uniform_choice(i, walks, adj_nodes, adj_offset, choices, degrees, nb_degrees)\n",
        "            else:\n",
        "                chosen_edges = self.choice_fct(i, walks, adj_nodes, adj_offset, choices, degrees, nb_degrees)\n",
        "\n",
        "            # update nodes\n",
        "            walks[i+1] = adj_nodes[chosen_edges]\n",
        "\n",
        "            # update edge features\n",
        "            edge_feat[i + 1] = x_edge[chosen_edges]\n",
        "\n",
        "            o = min(s, i+1)\n",
        "            prev = walks[i+1-o:i+1]\n",
        "\n",
        "            if self.compute_id:\n",
        "                # get local identity relation\n",
        "                id_enc[i+1, s-o:] = torch.eq(walks[i+1].view(1, w), prev)\n",
        "\n",
        "            if self.compute_adj:\n",
        "                # look up edges in the bit-wise adjacency encoding\n",
        "                cur_id = node_id[walks[i+1]]\n",
        "                cur_int = (cur_id // 63).view(1, -1, 1).repeat(o, 1, 1)\n",
        "                edges[i + 1, s - o:] = (torch.gather(adj_bits[prev], 2, cur_int).view(o,-1) >> (cur_id % 63).view(1,-1)) % 2 == 1\n",
        "\n",
        "        # permute walks into the correct shapes\n",
        "        data.walk_nodes = walks.permute(1, 0)\n",
        "\n",
        "        # combine id, adj and edge features\n",
        "        feat = [edge_feat.permute(1, 2, 0)]\n",
        "        if self.compute_id:\n",
        "            feat.append(torch._cast_Float(id_enc.permute(2, 1, 0)))\n",
        "        if self.compute_adj:\n",
        "            feat.append(torch._cast_Float(edges.permute(2, 1, 0))[:, :-1, :])\n",
        "        data.walk_x = torch.cat(feat, dim=1)\n",
        "\n",
        "        return data\n",
        "\n",
        "class VNUpdate(torch.nn.Module):\n",
        "    def __init__(self, dim, config):\n",
        "        \"\"\"\n",
        "        Intermediate update layer for the virtual node\n",
        "        :param dim: Dimension of the latent node embeddings\n",
        "        :param config: Python Dict with the configuration of the CRaWl network\n",
        "        \"\"\"\n",
        "        super(VNUpdate, self).__init__()\n",
        "\n",
        "        self.mlp = torch.nn.Sequential(torch.nn.Linear(dim, dim, bias=False),\n",
        "                                       torch.nn.BatchNorm1d(dim),\n",
        "                                       torch.nn.ReLU(),\n",
        "                                       torch.nn.Dropout(config['dropout']),\n",
        "                                       torch.nn.Linear(dim, dim, bias=False))\n",
        "\n",
        "    def forward(self, data):\n",
        "        x = scatter_sum(data.h, data.batch, dim=0)\n",
        "        if data.vn_h is not None:\n",
        "            x = x + data.vn_h\n",
        "        data.vn_h = self.mlp(x)\n",
        "        data.h = data.h + data.vn_h[data.batch]\n",
        "        return data\n",
        "\n",
        "\n",
        "class ConvModule(torch.nn.Module):\n",
        "    def __init__(self, conv_dim, w_feat_dim, dim_in, dim_out, kernel_size, config):\n",
        "        \"\"\"\n",
        "        :param conv_dim: Hidden dimension of the convolutions\n",
        "        :param w_feat_dim: feature dimension of the walk feature tensors (without the node features)\n",
        "        :param dim_in: Dimension of the latent node embedding used as input\n",
        "        :param dim_out: Dimension of updated latent node embedding\n",
        "        :param kernel_size: Kernel size of the convolutions\n",
        "        :param config: Python Dict with the configuration of the CRaWl network\n",
        "        \"\"\"\n",
        "        super(ConvModule, self).__init__()\n",
        "\n",
        "        self.rescale = not dim_out == dim_in\n",
        "        if self.rescale:\n",
        "            self.rescale_op = torch.nn.Linear(dim_in, dim_out, bias=False)\n",
        "\n",
        "        self.convs = torch.nn.Sequential(\n",
        "            torch.nn.Conv1d(dim_in + w_feat_dim, conv_dim, kernel_size, padding=0, bias=False),\n",
        "            torch.nn.BatchNorm1d(conv_dim),\n",
        "            torch.nn.ReLU(),\n",
        "            torch.nn.Conv1d(conv_dim, conv_dim, kernel_size, padding=0, bias=False),\n",
        "            torch.nn.BatchNorm1d(conv_dim),\n",
        "            torch.nn.ReLU())\n",
        "\n",
        "        self.out = torch.nn.Sequential(torch.nn.Linear(conv_dim, 2 * dim_out, bias=False),\n",
        "                                       torch.nn.BatchNorm1d(2 * dim_out),\n",
        "                                       torch.nn.ReLU(),\n",
        "                                       torch.nn.Dropout(config['dropout']),\n",
        "                                       torch.nn.Linear(2 * dim_out, dim_out, bias=False))\n",
        "\n",
        "    def forward(self, data):\n",
        "        # rescale for the residual connection\n",
        "        if self.rescale:\n",
        "            h_r = self.rescale_op(data.h)\n",
        "        else:\n",
        "            h_r = data.h\n",
        "\n",
        "        # build walk feature tensor\n",
        "        h_c = torch.cat([data.h[data.walk_nodes].transpose(2, 1), data.walk_x], dim=1)\n",
        "\n",
        "        # apply the cnn\n",
        "        h_c = self.convs(h_c)\n",
        "\n",
        "        # pool in walklet embeddings into nodes\n",
        "        wl = h_c.shape[0] * h_c.shape[2]\n",
        "        h = scatter_mean(h_c.transpose(2, 1).reshape(wl, -1), data.walk_nodes_flatt, dim=0, dim_size=data.num_nodes)\n",
        "\n",
        "        # update embeddings\n",
        "        h = self.out(h)\n",
        "        data.h = h + h_r\n",
        "        return data\n",
        "\n",
        "\n",
        "class CRaWl(torch.nn.Module):\n",
        "    def __init__(self, config ,node_feat_dim, edge_feat_dim, out_dim, node_feat_enc=None, edge_feat_enc=None):\n",
        "        \"\"\"\n",
        "        :param model_dir: Directory to store model in\n",
        "        :param config: Python Dict that specifies the configuration of the model\n",
        "        :param node_feat_dim: Dimension of the node features\n",
        "        :param edge_feat_dim: Dimension of the edge features\n",
        "        :param out_dim: Output dimension\n",
        "        :param node_feat_enc: Optional initial embedding of node features\n",
        "        :param edge_feat_enc: Optional initial embedding of edge features\n",
        "        \"\"\"\n",
        "        super(CRaWl, self).__init__()\n",
        "        #self.model_dir = model_dir\n",
        "        self.config = config\n",
        "        self.out_dim = out_dim\n",
        "        self.node_feat_enc = node_feat_enc\n",
        "        self.edge_feat_enc = edge_feat_enc\n",
        "        self.layers = config['layers']\n",
        "        self.hidden = config['hidden_dim']\n",
        "        self.kernel_size = config['kernel_size']\n",
        "        self.dropout = config['dropout']\n",
        "        self.residual = config['residual'] if 'residual' in config.keys() else True\n",
        "        self.pool_op = config['pool'] if 'pool' in config.keys() else 'mean'\n",
        "        self.vn = config['vn'] if 'vn' in config.keys() else False\n",
        "\n",
        "        self.border = 2 * ((self.kernel_size - 1) // 2)\n",
        "        self.tail = (self.kernel_size - 1) * 2\n",
        "\n",
        "        self.walker = Walker(config)\n",
        "\n",
        "        self.walk_dim = self.walker.struc_feat_dim + edge_feat_dim\n",
        "        self.conv_dim = config['conv_dim'] if 'conv_dim' in config.keys() else self.hidden\n",
        "\n",
        "        modules = []\n",
        "        for i in range(self.layers):\n",
        "            modules.append(ConvModule(self.conv_dim, self.walk_dim, node_feat_dim if i == 0 else self.hidden, self.hidden, self.kernel_size, config))\n",
        "            if self.vn and i < self.layers - 1:\n",
        "                modules.append(VNUpdate(self.hidden, config))\n",
        "\n",
        "        self.convs = torch.nn.Sequential(*modules)\n",
        "\n",
        "        self.node_out = torch.nn.Sequential(torch.nn.BatchNorm1d(self.hidden), torch.nn.ReLU())\n",
        "\n",
        "        if config['graph_out'] == 'linear':\n",
        "            self.linear_out = True\n",
        "            self.graph_out = torch.nn.Sequential(torch.nn.Linear(self.hidden, out_dim))\n",
        "        else:\n",
        "            self.linear_out = False\n",
        "            self.graph_out = torch.nn.Sequential(torch.nn.Linear(self.hidden, self.hidden),\n",
        "                                                 torch.nn.ReLU(),\n",
        "                                                 torch.nn.Linear(self.hidden, out_dim))\n",
        "\n",
        "        pytorch_total_params = sum(p.numel() for p in self.parameters() if p.requires_grad)\n",
        "        print(f'Number of paramters: {pytorch_total_params}')\n",
        "\n",
        "\n",
        "    def forward(self, data, walk_steps=None, walk_start_p=1.0):\n",
        "        # apply initial node feature encoding (optional)\n",
        "        x_n = data.x\n",
        "        if self.node_feat_enc is not None:\n",
        "            x_n = self.node_feat_enc(x_n)\n",
        "        data.h = x_n\n",
        "\n",
        "        # apply initial edge feature encoding (optional)\n",
        "        x_e = data.edge_attr\n",
        "        if self.edge_feat_enc is not None:\n",
        "            x_e = self.edge_feat_enc(x_e)\n",
        "\n",
        "        # compute walks\n",
        "        data = self.walker.sample_walks(data, x_e, steps=walk_steps, start_p=walk_start_p)\n",
        "\n",
        "        # pre-compute array of center nodes across all walks\n",
        "        data.walk_nodes_flatt = data.walk_nodes[:, self.border:-self.border].reshape(-1)\n",
        "\n",
        "        if self.vn:\n",
        "            data.vn_h = None\n",
        "\n",
        "        # apply convolutions\n",
        "        self.convs(data)\n",
        "\n",
        "        # pool node embeddings\n",
        "        data.h = self.node_out(data.h)\n",
        "        if self.pool_op == 'sum':\n",
        "            x = scatter_sum(data.h, data.batch, dim=0)\n",
        "        else:\n",
        "            x = scatter_mean(data.h, data.batch, dim=0)\n",
        "\n",
        "        y = self.graph_out(x)\n",
        "        return y\n",
        "\n",
        "from torch_geometric.datasets import TUDataset\n",
        "from torch_geometric.data import DataLoader\n",
        "\n",
        "def get_idx(dataset, split_dir, fold):\n",
        "    with open(os.path.join(split_dir, dataset, f'train_idx-{fold + 1}.txt'), 'r') as f:\n",
        "        train_idx = [int(i) for i in f]\n",
        "    with open(os.path.join(split_dir, dataset, f'test_idx-{fold + 1}.txt'), 'r') as f:\n",
        "        test_idx = [int(i) for i in f]\n",
        "    return train_idx, test_idx\n",
        "\n",
        "from torch_geometric.utils import degree\n",
        "\n",
        "def transform(g):\n",
        "    g.x = torch.log(degree(g.edge_index[0], g.num_nodes)).view(-1, 1)\n",
        "    g = preproc(g)\n",
        "    return g\n",
        "import math\n",
        "def _split_rand(  split_ratio=0.7, seed=0, shuffle=True):\n",
        "        num_entries = 1113\n",
        "        indices = list(range(num_entries))\n",
        "        np.random.seed(seed)\n",
        "        np.random.shuffle(indices)\n",
        "        split = int(math.floor(split_ratio * num_entries))\n",
        "        train_idx, valid_idx = indices[:split], indices[split:]\n",
        "\n",
        "        print(\n",
        "            \"train_set : test_set = %d : %d\",\n",
        "            len(train_idx), len(valid_idx))\n",
        "\n",
        "        return train_idx, valid_idx\n",
        "def get_split_data( config):\n",
        "    dataset = TUDataset(root = '', name = 'PROTEINS', transform=transform)\n",
        "    #dataset = data_name\n",
        "    train_idx, test_idx = _split_rand()\n",
        "    np.random.shuffle(train_idx)\n",
        "    np.random.shuffle(test_idx)\n",
        "    train_iter = CRaWlLoader(list(dataset[train_idx]), batch_size=50, shuffle=True)\n",
        "    val_iter = CRaWlLoader(list(dataset[test_idx]), batch_size=50 , shuffle= True)\n",
        "\n",
        "    return train_iter, val_iter\n",
        "\n",
        "#dataset = TUDataset(name = 'PROTEINS' , root  ='')\n",
        "#print(dataset[0])\n",
        "#torch.manual_seed(12345)\n",
        "#dataset = dataset.shuffle()\n",
        "config = {\n",
        "    \"name\": \"default\",\n",
        "    \"pool\": \"mean\",\n",
        "    \"steps\": 20,\n",
        "    \"win_size\": 8,\n",
        "    \"walk_mode\": \"nb\",\n",
        "    \"hidden_dim\": 256,\n",
        "    \"layers\": 5,\n",
        "    \"kernel_size\": 5,\n",
        "    \"graph_out\": \"mlp\",\n",
        "    \"dropout\": 0.5,\n",
        "}\n",
        "train_loader , test_loader  = get_split_data(config )\n",
        "#test_dataset = dataset[779:]\n",
        "\n",
        "#train_loader = CRaWlLoader(dataset = train_dataset , shuffle  = True , batch_size  = 16)\n",
        "#test_loader = CRaWlLoader(dataset = test_dataset , shuffle  = True , batch_size  = 16)\n",
        "\n",
        "\n",
        "\n",
        "model = CRaWl(config = config, node_feat_dim = 1, edge_feat_dim = 1,out_dim =  2)\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "\n",
        "def train():\n",
        "    model.train()\n",
        "\n",
        "    for data in train_loader:  # Iterate in batches over the training dataset.\n",
        "         #print(data)\n",
        "         out = model(data)  # Perform a single forward pass.\n",
        "         #print(out)\n",
        "         #print(out)\n",
        "         #print(data.y)\n",
        "         #pred = out.argmax(dim=1)\n",
        "         #pred = pred.view(-1)\n",
        "         #print(pred.T)\n",
        "         #pred = pred.unsqueeze(1)\n",
        "         #print(pred)\n",
        "         #y = data.y\n",
        "         #y = y.squeeze(1)\n",
        "         #y = y.view(-1)\n",
        "         #print(y.T)\n",
        "         y = data.y.squeeze(1)\n",
        "         #print(y)\n",
        "         loss = criterion(out, y)  # Compute the loss.\n",
        "         #print('loss ',loss)\n",
        "         loss.backward()  # Derive gradients.\n",
        "         optimizer.step()  # Update parameters based on gradients.\n",
        "         optimizer.zero_grad()  # Clear gradients.\n",
        "\n",
        "def test(loader):\n",
        "     model.eval()\n",
        "\n",
        "     correct = 0\n",
        "     total = 0\n",
        "     for data in loader:  # Iterate in batches over the training/test dataset.\n",
        "         out = model(data)  \n",
        "         _, pred = torch.max(out.data, 1)# Use the class with highest probability.\n",
        "         total += len(data.y)\n",
        "         #print(pred , data.y)\n",
        "         for x , y in zip(pred , data.y):\n",
        "           if x == y:\n",
        "             correct += 1\n",
        "\n",
        "         #correct += (pred == data.y).sum().item() # Check against ground-truth labels.\n",
        "     print(correct) \n",
        "     return  1.0*correct / total  # Derive ratio of correct predictions.\n",
        "\n",
        "\n",
        "for epoch in range(50):\n",
        "    train()\n",
        "    #train_acc = test(train_loader)\n",
        "    test_acc = test(test_loader)\n",
        "    print('test Acc', {test_acc})\n",
        "\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "train_set : test_set = %d : %d 779 334\n",
            "Number of paramters: 4440834\n",
            "139\n",
            "test Acc {0.4161676646706587}\n",
            "180\n",
            "test Acc {0.5389221556886228}\n",
            "233\n",
            "test Acc {0.6976047904191617}\n",
            "232\n",
            "test Acc {0.6946107784431138}\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}