07-custom-dataset-solution.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "07-custom-dataset.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "duBL5_7Rb2Za"
      },
      "source": [
        "# Custom Dataset\n",
        "\n",
        "You are provided a pair of files with data.\n",
        "\n",
        "* `node.csv` -- a TSV file containing the node ID and a node feature vector of size 100.\n",
        "```\n",
        "0\t0.38250,0.55505,0.32324,0.69098,0.97875,0.50953,0.59311,0.93023,0.52275,0.21924\n",
        "```\n",
        "* `edges.csv` -- a TSV file with pairs of node IDs that are connected.\n",
        "```\n",
        "0\t277\n",
        "0\t445\n",
        "```\n",
        "\n",
        "Running the cell under __Create Raw Data__ section will write out the raw files under `random/raw`.\n",
        "\n",
        "Construct a custom dataset that extends `torch_geometric.data.Dataset` similar to the CORA dataset, i.e. with only one graph in the dataset. The custom dataset should `process` the raw files and write out a `Data` object in the `random/processed`.\n",
        "\n",
        "Instantiate the custom dataset, you should see new files appear in the `random/processed` folder.\n",
        "\n",
        "Verify that the length of the dataset is 1 and the number of node features is 100.\n",
        "\n",
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "ZjKqWQMJKgD8",
        "outputId": "24be8a05-29f5-4541-e99a-7831f2fbc07a"
      },
      "source": [
        "import torch\n",
        "torch.__version__"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'1.9.0+cu111'"
            ]
          },
          "metadata": {},
          "execution_count": 1
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a3adiZrCKn91"
      },
      "source": [
        "%%capture\n",
        "!pip install -q torch-scatter -f https://pytorch-geometric.com/whl/torch-1.9.0+cu111.html\n",
        "!pip install -q torch-sparse -f https://pytorch-geometric.com/whl/torch-1.9.0+cu111.html\n",
        "!pip install -q torch-geometric"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SUH_tLX8b7RU"
      },
      "source": [
        "## Create Raw Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bfYCY0TOKqka"
      },
      "source": [
        "import numpy as np\n",
        "import os\n",
        "import shutil\n",
        "import torch_geometric.utils as pyg_utils\n",
        "\n",
        "raw_dir = \"random/raw\"\n",
        "\n",
        "shutil.rmtree(\"random\")\n",
        "os.makedirs(raw_dir)\n",
        "for i in range(50):\n",
        "  edge_index = pyg_utils.barabasi_albert_graph(100, 50)\n",
        "  node_features = torch.rand((100, 10), dtype=torch.float32)\n",
        "  fnode = open(os.path.join(raw_dir, \"node-{:d}.csv\".format(i)), \"w\")\n",
        "  for j in range(node_features.size(0)):\n",
        "    features = node_features[j].numpy().tolist()\n",
        "    features_str = \",\".join([\"{:.5f}\".format(feat) for feat in features])\n",
        "    fnode.write(\"{:d}\\t{:s}\\n\".format(j, features_str))\n",
        "  fnode.close()\n",
        "  fedge = open(os.path.join(raw_dir, \"edge-{:d}.csv\".format(i)), \"w\")\n",
        "  for j in range(edge_index.size(1)):\n",
        "    src, dst = edge_index[0, j], edge_index[1, j]\n",
        "    fedge.write(\"{:d}\\t{:d}\\n\".format(src, dst))\n",
        "  fedge.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1KYZEKcwcLJp"
      },
      "source": [
        "## Custom Dataset\n",
        "\n",
        "Extend the `torch_geometric.data.Dataset` object to create a PyG Dataset object.\n",
        "\n",
        "Use the code in the Pytorch documentation page [Creating your own Dataset](https://pytorch-geometric.readthedocs.io/en/latest/notes/create_dataset.html) and/or the [Youtube video on Creating a Custom Dataset in Pytorch Geometric](https://www.youtube.com/watch?v=QLIkOtKS4os) by DeepFindr as examples.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DW6Zi6X-MYwc"
      },
      "source": [
        "import torch_geometric.data as pyg_data"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pJFrlw__VK5T"
      },
      "source": [
        "# shutil.rmtree(\"random/processed\", ignore_errors=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rDJOosESMJ9a"
      },
      "source": [
        "class RandomDataset(pyg_data.Dataset):\n",
        "  def __init__(self, root, transform=None, pre_transform=None):\n",
        "    super().__init__(root, transform, pre_transform)\n",
        "\n",
        "  @property\n",
        "  def raw_file_names(self):\n",
        "    raw_files = []\n",
        "    for i in range(50):\n",
        "      raw_files.append(\"node-{:d}.csv\".format(i))\n",
        "      raw_files.append(\"edge-{:d}.csv\".format(i))\n",
        "    return raw_files\n",
        "\n",
        "  @property\n",
        "  def processed_file_names(self):\n",
        "    processed_files = []\n",
        "    for i in range(50):\n",
        "      processed_files.append(\"random-{:d}.pt\".format(i))\n",
        "    return processed_files\n",
        "\n",
        "  def download(self):\n",
        "    pass\n",
        "\n",
        "  def process(self):\n",
        "    for i in range(50):\n",
        "      features, edge_index = [], []\n",
        "      with open(os.path.join(self.raw_dir, \"node-{:d}.csv\".format(i)), \"r\") as fnode:\n",
        "        for line in fnode:\n",
        "          node_id, node_feats = line.strip().split('\\t')\n",
        "          features.append(np.array([float(x) for x in node_feats.split(',')]))\n",
        "      features = torch.tensor(np.array(features), dtype=torch.float32)\n",
        "      with open(os.path.join(self.raw_dir, \"edge-{:d}.csv\".format(i)), \"r\") as fedge:\n",
        "        for line in fedge:\n",
        "          src, dst = line.strip().split('\\t')\n",
        "          edge_index.append((int(src), int(dst)))\n",
        "          edge_index.append((int(dst), int(src)))\n",
        "      edge_index = torch.tensor(edge_index).t().to(torch.long)\n",
        "      labels = torch.randint(low=0, high=2, size=(1,))\n",
        "      data = pyg_data.Data(x=features,\n",
        "                          edge_index=edge_index, \n",
        "                          y=labels)\n",
        "      torch.save(data, os.path.join(\n",
        "          self.processed_dir, \"random-{:d}.pt\".format(i)))\n",
        "\n",
        "  def len(self):\n",
        "    return len(self.processed_file_names)\n",
        "\n",
        "  def get(self, idx):\n",
        "    data = torch.load(os.path.join(self.processed_dir, 'random-{:d}.pt'.format(idx)))\n",
        "    return data\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QpLOJREBqwH-"
      },
      "source": [
        "## Test Custom Dataset\n",
        "\n",
        "You should be able to instantiate your new custom Dataset and verify its properties.\n",
        "\n",
        "* `num_features` -- should be 10\n",
        "* `len()` -- should be 50"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NH_HdfemTTv8",
        "outputId": "1465c50f-fc9e-4c81-f15c-f5c7f4ee6c87"
      },
      "source": [
        "dataset = RandomDataset(\"random\")\n",
        "dataset"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Processing...\n",
            "Done!\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "RandomDataset(50)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QVNatdwPTZHn",
        "outputId": "7be5da8a-fb8b-43ca-8b30-950034593396"
      },
      "source": [
        "len(dataset)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "50"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-rJwsspKYv3C",
        "outputId": "bfa6e7e4-09ef-4fcc-f78c-bea3e0a097d5"
      },
      "source": [
        "dataset.num_features"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "10"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oyVuxcvVlw_c",
        "outputId": "61f70f22-0405-4d14-e56c-01eec39b4900"
      },
      "source": [
        "dataset[0]"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Data(x=[100, 10], edge_index=[2, 5536], y=[1])"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pft_XU4scUPI"
      },
      "source": [
        "## Wrap in DataLoader\n",
        "\n",
        "Try wrapping your custom Dataset into a PyG DataLoader and iterate through one batch. Print out the batch and verify that it looks right (for example, the number of labels should be the same as your batch size)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c18ox6ThZpe8",
        "outputId": "e3325fd9-a578-4505-d55c-13ce2f7c07b8"
      },
      "source": [
        "from torch_geometric.loader import DataLoader\n",
        "\n",
        "loader = DataLoader(dataset, batch_size=32, shuffle=True)\n",
        "for batch in loader:\n",
        "  print(batch)\n",
        "  break"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "DataBatch(x=[3200, 10], edge_index=[2, 173920], y=[32], batch=[3200], ptr=[33])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_Zba4CZQbe1C"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}