docs/sources.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import webdataset as wds\n",
    "import braceexpand\n",
    "from torch.utils.data import IterableDataset\n",
    "from webdataset import gopen"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Local and Remote Storage URLs\n",
    "\n",
    "WebDataset refers to data sources using file paths or URLs. The following are all valid ways of referring to a data source:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = wds.WebDataset(\"dataset-000.tar\")\n",
    "dataset = wds.WebDataset(\"file:dataset-000.tar\")\n",
    "dataset = wds.WebDataset(\"http://server/dataset-000.tar\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An additional way of referring to data is using the `pipe:` scheme, so the following is also equivalent to the above references:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = wds.WebDataset(\"pipe:cat dataset-000.tar\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use the same notation for accessing data in cloud storage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = wds.WebDataset(\"pipe:gsutil cat gs://somebucket/dataset-000.tar\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that access to standard web schemas are implemented using `curl`. That is, `http://server/dataset.tar` is internally simply treated like `pipe:curl -s -L 'http://server/dataset.tar'`. The use of `curl` to access Internet protocols actually is more efficient than using the built-in `http` library because it results in asynchronous name resolution and downloads."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "File opening is handled by `webdataset.gopen.gopen`. This is a small function that just wraps standard Python file I/O and pipe capabilities.\n",
    "\n",
    "You can define handlers for new schemes or override implementations for existing schemes by adding entries to `wds.gopen_schemes`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gopen_gs(url, mode=\"rb\", bufsize=8192):\n",
    "    ...\n",
    "\n",
    "gopen.gopen_schemes[\"gs\"] = gopen_gs "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Standard Input/Output\n",
    "\n",
    "For the following examples, assume that we have a program called `image-classifier` that takes a WebDataset containing just JPEG files as input and produces a WebDataset containing JPEG files and their corresponding classifications in JSON format:\n",
    "\n",
    "```Bash\n",
    "image-classifier input-shard.tar --output=output-shard.tar --model=some-model.pth\n",
    "```\n",
    "\n",
    "As a special case, the string \"-\" refers to standard input (reading) or standard output (writing). This allows code using WebDataset to be used as part of pipes. This is useful, for example, inside Kubernetes containers with limited local storage. Assume that you store shards in Google Cloud and access it with `gsutil`. Using \"-\", you can simply write:\n",
    "\n",
    "```Bash\n",
    "gsutil cat gs://input-bucket/data-000174.tar | image-classifer - -o - | gsutil cp - gs://output-bucket/output-000174.tar\n",
    "```\n",
    "\n",
    "It's also useful to create shards on the fly using `tar` and extract the result immediately; this lets you use shard based programs directly for operating on individual files. For example, for the `image-classifier` program above, you can write:\n",
    "\n",
    "```Bash\n",
    "tar cf - *.jpg | shard-classifier - -o - | tar xvf - --include '.json'\n",
    "```\n",
    "\n",
    "This is the rough equivalent of:\n",
    "\n",
    "```Bash\n",
    "for fname in *.jpg; do\n",
    "   image-classifier $fname > $(basename $fname .jpg).cls\n",
    "done\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multiple Shards and Mixing Datasets\n",
    "\n",
    "The `WebDataset` and `ShardList` classes take either a string or a list of strings as an argument. When given a string, the string is expanded using `braceexpand`. Therefore, the following three datasets are equivalent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = wds.WebDataset([\"dataset-000.tar\", \"dataset-001.tar\", \"dataset-002.tar\", \"dataset-003.tar\"])\n",
    "dataset = wds.WebDataset(\"dataset-{000..003}.tar\")\n",
    "dataset = wds.WebDataset(\"file:dataset-{000..003}.tar\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For complex training problems, you may want to mix multiple datasets, where each dataset consists of multiple shards. A good way is to expand each shard spec individually using `braceexpand` and concatenate the lists. Then you can pass the result list as an argument to `WebDataset`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1695\n"
     ]
    }
   ],
   "source": [
    "urls = (\n",
    "    list(braceexpand.braceexpand(\"imagenet-{000000..000146}.tar\")) +\n",
    "    list(braceexpand.braceexpand(\"openimages-{000000..000547}.tar\")) +\n",
    "    list(braceexpand.braceexpand(\"custom-images-{000000..000999}.tar\"))\n",
    ")\n",
    "print(len(urls))\n",
    "dataset = wds.WebDataset(urls, shardshuffle=True).shuffle(10000).decode(\"torchrgb\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mixing Datsets with a Custom `IterableDataset` Class"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For more complex sampling problems, you can also write sample processors. For example, to sample equally from several datasets, you could write something like this (the `Shorthands` and `Composable` base classes just add some convenience methods):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SampleEqually(IterableDataset, wds.Shorthands, wds.Composable):\n",
    "    def __init__(self, datasets):\n",
    "        super().__init__()\n",
    "        self.datasets = datasets\n",
    "    def __iter__(self):\n",
    "        sources = [iter(ds) for ds in self.datasets]\n",
    "        while True:\n",
    "            for source in sources:\n",
    "                try:\n",
    "                    yield next(source)\n",
    "                except StopIteration:\n",
    "                    return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can mix samples from different sources in more complex ways:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset1 = wds.WebDataset(\"imagenet-{000000..000146}.tar\", shardshuffle=True).shuffle(1000).decode(\"torchrgb\")\n",
    "dataset2 = wds.WebDataset(\"openimages-{000000..000547}.tar\", shardshuffle=True).shuffle(1000).decode(\"torchrgb\")\n",
    "dataset3 = wds.WebDataset(\"custom-images-{000000..000999}.tar\", shardshuffle=True).shuffle(1000).decode(\"torchrgb\")\n",
    "dataset = SampleEqually([dataset1, dataset2, dataset3]).shuffle(1000)"
   ]
  }
 ],
 "metadata": {
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}