Intro_to_Metadataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e5O1UdsY202_"
   },
   "source": [
    "##### Copyright 2019 Google LLC.\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Zy3bZKW82xP9"
   },
   "outputs": [],
   "source": [
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GXhLzrXN27af"
   },
   "source": [
    "# Using the Meta-Dataset Data Pipeline\n",
    "\n",
    "This notebook shows how to use `meta_dataset`’s input pipeline to sample data for the Meta-Dataset benchmark. There are two main ways in which data is sampled:\n",
    "1. **episodic**:  Returns N-way classification *episodes*, which contain a *support* (training) set and a *query* (test) set. The number of classes (N) may vary from episode to episode.\n",
    "2. **batch**:  Returns batches of images and their corresponding label, sampled from all available classes.\n",
    "\n",
    "We first import `meta_dataset` and other required packages, and define utility functions for visualization. We’ll make use of `meta_dataset.data.learning_spec` and `meta_dataset.data.pipeline`; their purpose will be made clear later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "cellView": "both",
    "colab": {},
    "colab_type": "code",
    "id": "ZyMqBhZIxPQD"
   },
   "outputs": [],
   "source": [
    "#@title Imports and Utility Functions\n",
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import os\n",
    "from collections import Counter\n",
    "import gin\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "from meta_dataset.data import config\n",
    "from meta_dataset.data import dataset_spec as dataset_spec_lib\n",
    "from meta_dataset.data import learning_spec\n",
    "from meta_dataset.data import pipeline\n",
    "\n",
    "\n",
    "def plot_episode(support_images, support_class_ids, query_images,\n",
    "                 query_class_ids, size_multiplier=1, max_imgs_per_col=10,\n",
    "                 max_imgs_per_row=10):\n",
    "  for name, images, class_ids in zip(('Support', 'Query'),\n",
    "                                     (support_images, query_images),\n",
    "                                     (support_class_ids, query_class_ids)):\n",
    "    n_samples_per_class = Counter(class_ids)\n",
    "    n_samples_per_class = {k: min(v, max_imgs_per_col)\n",
    "                           for k, v in n_samples_per_class.items()}\n",
    "    id_plot_index_map = {k: i for i, k\n",
    "                         in enumerate(n_samples_per_class.keys())}\n",
    "    num_classes = min(max_imgs_per_row, len(n_samples_per_class.keys()))\n",
    "    max_n_sample = max(n_samples_per_class.values())\n",
    "    figwidth = max_n_sample\n",
    "    figheight = num_classes\n",
    "    if name == 'Support':\n",
    "      print('#Classes: %d' % len(n_samples_per_class.keys()))\n",
    "    figsize = (figheight * size_multiplier, figwidth * size_multiplier)\n",
    "    fig, axarr = plt.subplots(\n",
    "        figwidth, figheight, figsize=figsize)\n",
    "    fig.suptitle('%s Set' % name, size='20')\n",
    "    fig.tight_layout(pad=3, w_pad=0.1, h_pad=0.1)\n",
    "    reverse_id_map = {v: k for k, v in id_plot_index_map.items()}\n",
    "    for i, ax in enumerate(axarr.flat):\n",
    "      ax.patch.set_alpha(0)\n",
    "      # Print the class ids, this is needed since, we want to set the x axis\n",
    "      # even there is no picture.\n",
    "      ax.set(xlabel=reverse_id_map[i % figheight], xticks=[], yticks=[])\n",
    "      ax.label_outer()\n",
    "    for image, class_id in zip(images, class_ids):\n",
    "      # First decrement by one to find last spot for the class id.\n",
    "      n_samples_per_class[class_id] -= 1\n",
    "      # If class column is filled or not represented: pass.\n",
    "      if (n_samples_per_class[class_id] < 0 or\n",
    "          id_plot_index_map[class_id] >= max_imgs_per_row):\n",
    "        continue\n",
    "      # If width or height is 1, then axarr is a vector.\n",
    "      if axarr.ndim == 1:\n",
    "        ax = axarr[n_samples_per_class[class_id]\n",
    "                   if figheight == 1 else id_plot_index_map[class_id]]\n",
    "      else:\n",
    "        ax = axarr[n_samples_per_class[class_id], id_plot_index_map[class_id]]\n",
    "      ax.imshow(image / 2 + 0.5)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def plot_batch(images, labels, size_multiplier=1):\n",
    "  num_examples = len(labels)\n",
    "  figwidth = np.ceil(np.sqrt(num_examples)).astype('int32')\n",
    "  figheight = num_examples // figwidth\n",
    "  figsize = (figwidth * size_multiplier, (figheight + 1.5) * size_multiplier)\n",
    "  _, axarr = plt.subplots(figwidth, figheight, dpi=300, figsize=figsize)\n",
    "\n",
    "  for i, ax in enumerate(axarr.transpose().ravel()):\n",
    "    # Images are between -1 and 1.\n",
    "    ax.imshow(images[i] / 2 + 0.5)\n",
    "    ax.set(xlabel=labels[i], xticks=[], yticks=[])\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BOn_YZdqPIv5"
   },
   "source": [
    "# Primers\n",
    "1. Download your data and process it as explained in [link](https://github.com/google-research/meta-dataset/blob/main/README.md#downloading-and-converting-datasets). Set `BASE_PATH` pointing the processed tf-records (`$RECORDS` in the conversion instructions).\n",
    "2. `meta_dataset` supports many different setting for sampling data. We use [gin-config](https://github.com/google/gin-config) to control default parameters of our functions. You can go to default gin file we are pointing and see the default values.\n",
    "3. You can use `meta_dataset` in **eager** or **graph** mode.\n",
    "4. Let's write a generator that makes the right calls to return data from dataset. `dataset.make_one_shot_iterator()` returns an iterator where each element is an episode.\n",
    "4. SPLIT is used to define which part of the meta-split is going to be used. Different splits have different classes and the details on how they are created can be found in the [paper](https://arxiv.org/abs/1903.03096)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_di9Tczj8joM"
   },
   "outputs": [],
   "source": [
    "# 1\n",
    "BASE_PATH = '/path/to/records'\n",
    "GIN_FILE_PATH = 'meta_dataset/learn/gin/setups/data_config.gin'\n",
    "# 2\n",
    "gin.parse_config_file(GIN_FILE_PATH)\n",
    "# 3\n",
    "# Comment out to disable eager execution.\n",
    "tf.enable_eager_execution()\n",
    "# 4\n",
    "def iterate_dataset(dataset, n):\n",
    "  if not tf.executing_eagerly():\n",
    "    iterator = dataset.make_one_shot_iterator()\n",
    "    next_element = iterator.get_next()\n",
    "    with tf.Session() as sess:\n",
    "      for idx in range(n):\n",
    "        yield idx, sess.run(next_element)\n",
    "  else:\n",
    "    for idx, episode in enumerate(dataset):\n",
    "      if idx == n:\n",
    "        break\n",
    "      yield idx, episode\n",
    "# 5\n",
    "SPLIT = learning_spec.Split.TRAIN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Pn6ndPMhxs8W"
   },
   "source": [
    "# Reading datasets\n",
    "In order to sample data, we need to read the dataset_spec files for each dataset. Following snippet reads those files into a list. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Z0uU6WrbxsMa"
   },
   "outputs": [],
   "source": [
    "ALL_DATASETS = ['aircraft', 'cu_birds', 'dtd', 'fungi', 'ilsvrc_2012',\n",
    "                'omniglot', 'quickdraw', 'vgg_flower']\n",
    "\n",
    "all_dataset_specs = []\n",
    "for dataset_name in ALL_DATASETS:\n",
    "  dataset_records_path = os.path.join(BASE_PATH, dataset_name)\n",
    "  dataset_spec = dataset_spec_lib.load_dataset_spec(dataset_records_path)\n",
    "  all_dataset_specs.append(dataset_spec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7p448EXYxwbb"
   },
   "source": [
    "# (1) Episodic Mode\n",
    "`meta_dataset` uses [tf.data.Dataset](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) API and it takes one call to `pipeline.make_multisource_episode_pipeline()`. We loaded or defined most of the variables used during this call above. The remaining parameters are explained below:\n",
    "\n",
    "- **use_bilevel_ontology_list**:  This is a list of booleans indicating whether corresponding dataset in `ALL_DATASETS` should use bilevel ontology. Omniglot is set up with a hierarchy with two level: the alphabet (Latin, Inuktitut...), and the character (with 20 examples per character).\n",
    "The flag means that each episode will contain classes from a single alphabet. \n",
    "- **use_dag_ontology_list**:  This is a list of booleans indicating whether corresponding dataset in `ALL_DATASETS` should use dag_ontology. Same idea for ImageNet, except it uses the hierarchical sampling procedure described in the article.\n",
    "- **image_size**: All images from various datasets are down or upsampled to the same size. This is the flag controls the edge size of the square.\n",
    "- **shuffle_buffer_size**: Controls the amount of shuffling among examples from any given class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jPlnBWwkwuGP"
   },
   "outputs": [],
   "source": [
    "use_bilevel_ontology_list = [False]*len(ALL_DATASETS)\n",
    "use_dag_ontology_list = [False]*len(ALL_DATASETS)\n",
    "# Enable ontology aware sampling for Omniglot and ImageNet. \n",
    "use_bilevel_ontology_list[5] = True\n",
    "use_dag_ontology_list[4] = True\n",
    "variable_ways_shots = config.EpisodeDescriptionConfig(\n",
    "    num_query=None, num_support=None, num_ways=None)\n",
    "\n",
    "dataset_episodic = pipeline.make_multisource_episode_pipeline(\n",
    "    dataset_spec_list=all_dataset_specs,\n",
    "    use_dag_ontology_list=use_dag_ontology_list,\n",
    "    use_bilevel_ontology_list=use_bilevel_ontology_list,\n",
    "    episode_descr_config=variable_ways_shots,\n",
    "    split=SPLIT,\n",
    "    image_size=84,\n",
    "    shuffle_buffer_size=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BN66UXO79Bo2"
   },
   "source": [
    "## Using Dataset\n",
    "1. The episodic dataset consist in a tuple of the form (Episode, data source ID). The data source ID is an integer Tensor containing a value in the range [0, len(all_dataset_specs) - 1]\n",
    "signifying which of the datasets of the multisource pipeline the given episode\n",
    "came from. Episodes consist of support and query sets and we want to learn to classify images at the query set correctly given the support images. For both support and query set we have `images`, `labels` and `class_ids`. Labels are transformed class_ids offset to zero, so that global class_ids are set to \\[0, N\\] where N is the number of classes in an episode.\n",
    "3. As one can see the number of images in query set and support set is different. Images are scaled, copied into 84\\*84\\*3 tensors. Labels are presented in two forms:\n",
    "   * `*_labels` are relative to the classes selected for the current episode only. They are used as targets for this episode.\n",
    "   * `*_class_ids` are the original class ids relative to the whole dataset. They are used for visualization and diagnostics.\n",
    "4. It easy to convert tensors of the episode into numpy arrays and use them outside of the Tensorflow framework.\n",
    "5. Classes might have different number of samples in the support set, whereas each class has 10 samples in the query set. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lomtjv9rw5WP"
   },
   "outputs": [],
   "source": [
    "# 1\n",
    "idx, (episode, source_id) = next(iterate_dataset(dataset_episodic, 1))\n",
    "print('Got an episode from dataset:', all_dataset_specs[source_id].name)\n",
    "\n",
    "# 2\n",
    "for t, name in zip(episode,\n",
    "                   ['support_images', 'support_labels', 'support_class_ids',\n",
    "                    'query_images', 'query_labels', 'query_class_ids']):\n",
    "  print(name, t.shape)\n",
    "\n",
    "# 3\n",
    "episode = [a.numpy() for a in episode]\n",
    "\n",
    "# 4\n",
    "support_class_ids, query_class_ids = episode[2], episode[5]\n",
    "print(Counter(support_class_ids))\n",
    "print(Counter(query_class_ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KxdVUqJiWmTX"
   },
   "source": [
    "## Visualizing Episodes\n",
    "Let's visualize the episodes. \n",
    "\n",
    "- Support and query set for each episode plotted sequentially. Set N_EPISODES to control number of episodes visualized.\n",
    "- Each episode is sampled from a single dataset and include N different classes. Each class might have different number of samples in support set, whereas number of images in query set is fixed. We limit number of classes and images per class to 10 in order to create legible plots. Actual episodes might have more classes and samples.  \n",
    "- Each column represents a distinct class and dataset specific class ids are plotted on the x_axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9v2ePLTkoZlE"
   },
   "outputs": [],
   "source": [
    "# 1\n",
    "N_EPISODES=2\n",
    "# 2, 3\n",
    "for idx, (episode, source_id) in iterate_dataset(dataset_episodic, N_EPISODES):\n",
    "  print('Episode id: %d from source %s' % (idx, all_dataset_specs[source_id].name))\n",
    "  episode = [a.numpy() for a in episode]\n",
    "  plot_episode(support_images=episode[0], support_class_ids=episode[2],\n",
    "               query_images=episode[3], query_class_ids=episode[5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pL2AZ5gx3cDS"
   },
   "source": [
    "# (2) Batch Mode\n",
    "Second mode that `meta_dataset` library provides is the batch mode, where one can sample batches from the list of  datasets in a non-episodic manner and use it to train baseline models. There are couple things to note here:\n",
    "\n",
    "- Each batch is sampled from a different dataset.\n",
    "- `ADD_DATASET_OFFSET` controls whether the class_id's returned by the iterator overlaps among different datasets or not. A dataset specific offset is added in order to make returned ids unique.\n",
    "- `make_multisource_batch_pipeline()` creates a `tf.data.Dataset` object that returns datasets of the form (Batch, data source ID) where similarly to the\n",
    "episodic case, the data source ID is an integer Tensor that identifies which\n",
    "dataset the given batch originates from.\n",
    "- `shuffle_buffer_size` controls the amount of shuffling done among examples from a given dataset (unlike for the episodic pipeline)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jYY5zd_S6uG6"
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "ADD_DATASET_OFFSET = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BgkTLKKXPh8M"
   },
   "outputs": [],
   "source": [
    "dataset_batch = pipeline.make_multisource_batch_pipeline(\n",
    "    dataset_spec_list=all_dataset_specs, batch_size=BATCH_SIZE, split=SPLIT,\n",
    "    image_size=84, add_dataset_offset=ADD_DATASET_OFFSET,\n",
    "    shuffle_buffer_size=1000)\n",
    "\n",
    "for idx, ((images, labels), source_id) in iterate_dataset(dataset_batch, 1):\n",
    "  print(images.shape, labels.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7hGjt6GGonAz"
   },
   "outputs": [],
   "source": [
    "N_BATCH = 2\n",
    "for idx, (batch, source_id) in iterate_dataset(dataset_batch, N_BATCH):\n",
    "  print('Batch-%d from source %s' % (idx, all_dataset_specs[source_id].name))\n",
    "  plot_batch(*map(lambda a: a.numpy(), batch), size_multiplier=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tu4-jz89xt1f"
   },
   "source": [
    "# (3) Fixing Ways and Shots\n",
    "1. `meta_dataset` library provides option to set number of classes/samples per episode. There are 3 main flags you can set. \n",
    "    - **NUM_WAYS**: Fixes the # classes per episode. We would still get variable number of samples per class in the support set.\n",
    "    - **NUM_SUPPORT**: Fixes # samples per class in the support set.\n",
    "    - **NUM_SUPPORT**: Fixes # samples per class in the query set.\n",
    "2. If we want to use fixed `num_ways`, we have to disable ontology based sampling for omniglot and imagenet. We advise using single dataset for using this feature, since using multiple datasets is not supported/tested. In this notebook, we are using Quick, Draw! Dataset.\n",
    "3. We sample episodes and visualize them as we did earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8raM-sad6Igu"
   },
   "outputs": [],
   "source": [
    "#1\n",
    "NUM_WAYS = 8\n",
    "NUM_SUPPORT = 3\n",
    "NUM_QUERY = 5\n",
    "fixed_ways_shots = config.EpisodeDescriptionConfig(\n",
    "    num_ways=NUM_WAYS, num_support=NUM_SUPPORT, num_query=NUM_QUERY)\n",
    "\n",
    "#2\n",
    "use_bilevel_ontology_list = [False]*len(ALL_DATASETS)\n",
    "use_dag_ontology_list = [False]*len(ALL_DATASETS)\n",
    "quickdraw_spec = [all_dataset_specs[6]]\n",
    "#3\n",
    "dataset_fixed = pipeline.make_multisource_episode_pipeline(\n",
    "    dataset_spec_list=quickdraw_spec, use_dag_ontology_list=[False],\n",
    "    use_bilevel_ontology_list=use_bilevel_ontology_list, split=SPLIT,\n",
    "    image_size=84, episode_descr_config=fixed_ways_shots)\n",
    "\n",
    "N_EPISODES = 2\n",
    "for idx, (episode, source_id) in iterate_dataset(dataset_fixed, N_EPISODES):\n",
    "  print('Episode id: %d from source %s' % (idx, quickdraw_spec[source_id].name))\n",
    "  episode = [a.numpy() for a in episode]\n",
    "  plot_episode(support_images=episode[0], support_class_ids=episode[2],\n",
    "               query_images=episode[3], query_class_ids=episode[5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4tKDA6JFxt11"
   },
   "source": [
    "# (4) Using Meta-dataset with PyTorch\n",
    "As mentioned above it is super easy to consume `meta_dataset` as NumPy arrays. This also enables easy integration into other popular deep learning frameworks like PyTorch. TensorFlow code processes the data and passes it to PyTorch, ready to be consumed. Since the data loader and processing steps do not have any operation on the GPU, TF should not attempt to grab the GPU, and it should be available for PyTorch.\n",
    "1. Let's use an episodic dataset created earlier, `dataset_episodic`, and build on top of it. We will transpose tensor to CHW, which is the common order used by [convolutional layers](https://pytorch.org/docs/stable/nn.html?highlight=conv2d#torch.nn.functional.conv2d) of PyTorch. \n",
    "2. We will use zero-indexed labels, therefore grabbing `e[1]` and `e[4]`. At the end we return a generator that consumes the `tf.Dataset`. \n",
    "3. Using `.cuda()` on PyTorch tensors should distribute them to appropriate devices.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2d5w2YW-xt14"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "# 1\n",
    "to_torch_labels = lambda a: torch.from_numpy(a.numpy()).long()\n",
    "to_torch_imgs = lambda a: torch.from_numpy(np.transpose(a.numpy(), (0, 3, 1, 2)))\n",
    "# 2\n",
    "def data_loader(n_batches):\n",
    "  for i, (e, _) in enumerate(dataset_episodic):\n",
    "    if i == n_batches:\n",
    "      break\n",
    "    yield (to_torch_imgs(e[0]), to_torch_labels(e[1]),\n",
    "           to_torch_imgs(e[3]), to_torch_labels(e[4]))\n",
    "\n",
    "for i, batch in enumerate(data_loader(n_batches=2)):\n",
    "  #3\n",
    "  data_support, labels_support, data_query, labels_query = [x.cuda() for x in batch]\n",
    "  print(data_support.shape, labels_support.shape, data_query.shape, labels_query.shape) "
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "e5O1UdsY202_"
   ],
   "last_runtime": {
    "build_target": "//learning/brain/python/client:colab_notebook_py3",
    "kind": "private"
   },
   "name": "Intro to Metadataset.ipynb",
   "provenance": [
    {
     "file_id": "1z83txZ92A3930dqYb8kw-NF8IUhQIHk0",
     "timestamp": 1554403615924
    }
   ]
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}