07_Building_efficient_custom_data_loaders.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "144b1420-f3d3-488c-be92-93d7642cc0ce",
   "metadata": {},
   "source": [
    "To make your code run with maximum efficiency you also need to load your data efficiently into your device's memory. Fortunately PyTorch offers a tool to make data loading easy. It's called a DataLoader. A DataLoader uses multiple workers to simultanously load data from a Dataset and optionally uses a Sampler to sample data entries and form a batch.\n",
    "\n",
    "If you can randomly access your data, using a DataLoader is very easy: You simply need to implement a Dataset class that implements `__getitem__` (to read each data item) and `__len__` (to return the number of items in the dataset) methods. For example here's how to load images from a given directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "79d4b17c-730b-4ed6-aaaf-b40985658033",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "class CustomTensorDataset(Dataset):\n",
    "    def __init__(self, data_tensor, target_tensor, transform=None):\n",
    "        self.data_tensor = data_tensor\n",
    "        self.target_tensor = target_tensor\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data_tensor)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = self.data_tensor[idx]\n",
    "        target = self.target_tensor[idx]\n",
    "        \n",
    "        if self.transform:\n",
    "            data = self.transform(data)\n",
    "\n",
    "        return data, target\n",
    "\n",
    "# Let's create some random data tensors for demonstration\n",
    "data_tensor = torch.randn(100, 3, 32, 32)  # Assuming 100 images with shape (3, 32, 32)\n",
    "target_tensor = torch.randint(0, 10, (100,))  # Assuming 100 labels\n",
    "\n",
    "# Create an instance of the custom dataset\n",
    "dataset = CustomTensorDataset(data_tensor=data_tensor, target_tensor=target_tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83941aa2-bc6c-4da5-8718-00bfd671c382",
   "metadata": {},
   "source": [
    "You can then do the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5b92874d-fae9-4474-ac66-1ba02c8b49f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataloader = torch.utils.data.DataLoader(dataset, num_workers=2)\n",
    "# for data in dataloader:\n",
    "#     print(data[0].shape, data[1].shape)\n",
    "#     break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8a4305e-789c-4242-9fa2-75d5eae54d3f",
   "metadata": {},
   "source": [
    "Using a DataLoader to read data with random access may be ok if you have fast storage or if your data items are large. But imagine having a network file system with slow connection. Requesting individual files this way can be extremely slow and would probably end up becoming the bottleneck of your training pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c11a2dc-e651-4e45-9980-a67f23222381",
   "metadata": {},
   "source": [
    "A better approach is to store your data in a contiguous file format which can be read sequentially. For example if you have a large collection of images you can use tar to create a single archive and extract files from the archive sequentially in python. To do this you can use PyTorch's IterableDataset. To create an IterableDataset class you only need to implement an __iter__ method which sequentially reads and yields data items from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "58070881-a6ae-4ed1-bfc1-978c88638b73",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tarfile\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "\n",
    "from torch.utils.data import IterableDataset\n",
    "from torchvision.datasets.utils import extract_archive\n",
    "\n",
    "class TarImageDataset(IterableDataset):\n",
    "    def __init__(self, tar_path, transform=None):\n",
    "        self.tar_path = tar_path\n",
    "        self.transform = transform\n",
    "\n",
    "        # Extract the tar archive\n",
    "        self.extracted_folder = self.extract_tar()\n",
    "\n",
    "        # List all files in the extracted folder\n",
    "        self.file_list = os.listdir(self.extracted_folder)\n",
    "\n",
    "    def extract_tar(self):\n",
    "        # Extract the tar archive to a temporary folder\n",
    "        extracted_folder = tarfile.open(self.tar_path)\n",
    "        extracted_folder.extractall()\n",
    "        extracted_folder.close()\n",
    "        return os.path.splitext(self.tar_path)[0]  # Temporary folder name\n",
    "\n",
    "    def __iter__(self):\n",
    "        worker_info = torch.utils.data.get_worker_info()\n",
    "        if worker_info is None:  # In single-process data loading\n",
    "            return self.generate_samples()\n",
    "        else:  # In multi-process data loading\n",
    "            per_worker = int(len(self.file_list) / worker_info.num_workers)\n",
    "            worker_id = worker_info.id\n",
    "            start = worker_id * per_worker\n",
    "            end = (worker_id + 1) * per_worker\n",
    "            return iter(self.file_list[start:end])\n",
    "\n",
    "    def generate_samples(self):\n",
    "        for file_name in self.file_list:\n",
    "            image_path = os.path.join(self.extracted_folder, file_name)\n",
    "            image = Image.open(image_path).convert('RGB')\n",
    "            \n",
    "            if self.transform:\n",
    "                image = self.transform(image)\n",
    "\n",
    "            yield image\n",
    "\n",
    "# dataset = TarImageDataset(tar_path='images.tar', transform=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0f9af18-1df8-43ca-a8de-7b093bd9a30f",
   "metadata": {},
   "source": [
    "But there's a major problem with this implementation. If you try to use DataLoader to read from this dataset with more than one worker you'd observe a lot of duplicated images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "accc4444-00e7-48e6-96b9-b0879a2459c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataloader = torch.utils.data.DataLoader(TarImageDataset(\"/data/imagenet.tar\"), num_workers=8)\n",
    "# for data in dataloader:\n",
    "#     # data contains duplicated items"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60cd62eb-3792-4e97-bdcd-7942304e7c0a",
   "metadata": {},
   "source": [
    "The problem is that each worker creates a separate instance of the dataset and each would start from the beginning of the dataset. One way to avoid this is to instead of having one tar file, split your data into num_workers separate tar files and load each with a separate worker:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "95398f43-1e9f-4a25-b710-22c4ab4f72a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TarImageDataset(IterableDataset):\r\n",
    "    def __init__(self, paths):\r\n",
    "        super().__init__()\r\n",
    "        self.paths = paths\r\n",
    "\r\n",
    "    def __iter__(self):\r\n",
    "        worker_info = torch.utils.data.get_worker_info()\r\n",
    "        # For simplicity we assume num_workers is equal to number of tar files\r\n",
    "        if worker_info is None or worker_info.num_workers != len(self.paths):\r\n",
    "            raise ValueError(\"Number of workers doesn't match number of files.\")\r\n",
    "        yield from tar_image_iterator(self.paths[worker_info.worker_id])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac432e83-3a1d-4469-9f28-6322836d64ef",
   "metadata": {},
   "source": [
    "This is how our dataset class can be used:\r\n",
    "\r\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "aec0d2de-671d-4317-bc89-454f9bd1e101",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataloader = torch.utils.data.DataLoader(\n",
    "#     TarImageDataset([\"/data/imagenet_part1.tar\", \"/data/imagenet_part2.tar\"]), num_workers=2)\n",
    "# for data in dataloader:\n",
    "#     # do something with data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f2cbc1c-af62-4b8d-a4b6-f7e0e503cafc",
   "metadata": {},
   "source": [
    "We discussed a simple strategy to avoid duplicated entries problem. [tfrecord](https://github.com/vahidk/tfrecord) package uses slightly more sophisticated strategies to shard your data on the fly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5fb39ae-2670-4ab1-8264-0664abcd181d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}