pytorch_dataloader.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ac7d218",
   "metadata": {},
   "source": [
    "# Infinite dataloader\n",
    "This demonstrates an inifinte dataloader in pytorch, that loads data from a custom dataset.\n",
    "Th custom dataset is just for demo purpose. The torchvision.datasets.ImageFolder can be instead \n",
    "used to create the same.\n",
    "\n",
    "The folder structure is:\n",
    "\n",
    "/path/to/data/\n",
    "\n",
    "|---train/\\<class_folders\\>/\\<image_file\\>\n",
    "\n",
    "|---test/\\<class_folders\\>/\\<image_file\\>\n",
    "\n",
    "|---val/\\<class_folders\\>/\\<image_file\\>\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "11cca0e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import math\n",
    "import numpy as np\n",
    "import os\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "from PIL import Image\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import datasets, transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d5b4b294",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train path ../datasets/intel_image_classification/seg_train/seg_train \n",
      "test path ../datasets/intel_image_classification/seg_test/seg_test\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(True, True)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# setups the dataset paths\n",
    "dataset_path = Path('../datasets/intel_image_classification/')\n",
    "\n",
    "# Setup train and test paths\n",
    "train_dir = dataset_path / \"seg_train/seg_train\"\n",
    "test_dir = dataset_path / \"seg_test/seg_test\"\n",
    "\n",
    "print('train path {} \\ntest path {}'.format(train_dir, test_dir))\n",
    "\n",
    "os.path.isdir(train_dir), os.path.isdir(test_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "497d1a60",
   "metadata": {},
   "source": [
    "## Stanford format for preparing Dataset Class\n",
    "\n",
    "Here the stanford format for preparing custom dataset. A function is used to prepare two dictionaries segregating train-test set and mapping sample ids to labels\n",
    "\n",
    "\\>>> partition\n",
    "\n",
    "{'train': ['id-1', 'id-2', 'id-3'], 'validation': ['id-4']}\n",
    "\n",
    "\\>>> labels\n",
    "\n",
    "{'id-1': 0, 'id-2': 1, 'id-3': 2, 'id-4': 1}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "282de008",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dataset_preprocess(PATH) -> list:\n",
    "    \"\"\"\n",
    "    Given the following folder structure:\n",
    "    /path/to/data/\n",
    "    |---train/\\<class_folders\\>/\\<image_file\\>\n",
    "    |---test/\\<class_folders\\>/\\<image_file\\>\n",
    "    |---val/\\<class_folders\\>/\\<image_file\\>\n",
    "    \n",
    "    the function returns three dicts:\n",
    "    \n",
    "    1) partition: contains two dicts 'train' and 'validation' containing respective sample ids\n",
    "    2) maping of ids to labels\n",
    "    3) class name to ids\n",
    "    \n",
    "    Args:\n",
    "        param1 (str): 'path/to/data'\n",
    "    Returns:\n",
    "        list: [classes_to_idx, partition, labels, list of dicts]    \n",
    "    \"\"\"\n",
    "    \n",
    "    labels = {}\n",
    "    classes_to_idx = {}\n",
    "    partition = {'train': [], 'test': []}\n",
    "    \n",
    "    # print('PATH', PATH)\n",
    "    for file in glob.glob(os.path.join(PATH, '*/*/*/*.jpg')):\n",
    "        \n",
    "        # print('file: ', file)\n",
    "        \n",
    "        class_name = file.split('/')[-2]\n",
    "        # print('class_name', class_name)\n",
    "        if not class_name in classes_to_idx.keys():\n",
    "            classes_to_idx[class_name] = len(classes_to_idx.keys())\n",
    "        # print('classes_to_idx', classes_to_idx)\n",
    "        \n",
    "        ID = file\n",
    "        labels[ID] = classes_to_idx[class_name]\n",
    "\n",
    "        if 'seg_train' in file:\n",
    "            partition['train'].append(ID)\n",
    "        elif 'seg_test' in file:\n",
    "            partition['test'].append(ID)\n",
    "                \n",
    "        # print('partition', partition)\n",
    "        # print('labels', labels)\n",
    "        # break\n",
    "        \n",
    "    # print('classes_to_idx', classes_to_idx)\n",
    "    # print('partition train test len', len(partition['train']), len(partition['test']))\n",
    "    # print('labels', len(labels))\n",
    "    \n",
    "    return classes_to_idx, partition, labels\n",
    "    \n",
    "            \n",
    "_, _, _ = dataset_preprocess(dataset_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e7e6072b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extention of standard pytorch Dataset class. Returns X,y: images, labels after tranforms\n",
    "class MyDataset(Dataset):\n",
    "    'Characterizes a dataset for PyTorch'\n",
    "    def __init__(self, list_IDs, labels, classes_to_idx):\n",
    "        'Initialization'\n",
    "        self.labels = labels\n",
    "        self.list_IDs = list_IDs\n",
    "        self.classes_to_idx = classes_to_idx\n",
    "        self.transform = transforms.Compose([transforms.ToTensor(), transforms.Resize((224, 224))])\n",
    "    \n",
    "    def find_classes(self) -> dict:\n",
    "        'returns class names and indexes'\n",
    "        return self.classes_to_idx\n",
    "            \n",
    "    @staticmethod\n",
    "    def pil_loader(path: str) -> Image.Image:\n",
    "        'loads an image file'\n",
    "        with open(path, \"rb\") as f:\n",
    "            img = Image.open(f)\n",
    "            return img.convert(\"RGB\")\n",
    "\n",
    "    def __len__(self):\n",
    "        'Denotes the total number of samples'\n",
    "        return len(self.list_IDs)\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        'Generates one sample of data'\n",
    "        # Select sample\n",
    "        ID = self.list_IDs[index]\n",
    "\n",
    "        # Load data and get label\n",
    "        # print('ID', ID)\n",
    "        X = self.pil_loader(ID)\n",
    "        X = self.transform(X)\n",
    "        y = self.labels[ID]\n",
    "\n",
    "        return X, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d0eed2cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "class names {'sea': 0, 'mountain': 1, 'buildings': 2, 'glacier': 3, 'street': 4, 'forest': 5}\n",
      "dataset length 3000\n",
      "expected iterations 375\n",
      "torch.Size([8, 3, 224, 224]) 374 \r"
     ]
    },
    {
     "ename": "StopIteration",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mStopIteration\u001b[0m                             Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 19\u001b[0m\n\u001b[1;32m     17\u001b[0m idx \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m---> 19\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;28;43mnext\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43miter_obj\u001b[49m\u001b[43m)\u001b[49m[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mshape, idx, \u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\r\u001b[39;00m\u001b[38;5;124m'\u001b[39m, end\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     20\u001b[0m     idx\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/site-packages/torch/utils/data/dataloader.py:628\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    625\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    626\u001b[0m     \u001b[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[39;00m\n\u001b[1;32m    627\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()  \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[0;32m--> 628\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_next_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    629\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    630\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    631\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    632\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called:\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1306\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1304\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_persistent_workers:\n\u001b[1;32m   1305\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_shutdown_workers()\n\u001b[0;32m-> 1306\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m\n\u001b[1;32m   1308\u001b[0m \u001b[38;5;66;03m# Now `self._rcvd_idx` is the batch index we want to fetch\u001b[39;00m\n\u001b[1;32m   1309\u001b[0m \n\u001b[1;32m   1310\u001b[0m \u001b[38;5;66;03m# Check if the next sample has already been generated\u001b[39;00m\n\u001b[1;32m   1311\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_task_info[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_rcvd_idx]) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m2\u001b[39m:\n",
      "\u001b[0;31mStopIteration\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# test the standard pytroch dataloader which is expected to throw StopIteration exception\n",
    "classes_to_idx, partition, labels = dataset_preprocess(dataset_path)\n",
    "test_dataset = MyDataset(partition['test'], labels, classes_to_idx)\n",
    "\n",
    "print('class names', test_dataset.classes_to_idx)\n",
    "print('dataset length', len(test_dataset))\n",
    "\n",
    "BATCH_SIZE = 8\n",
    "train_dataloader = DataLoader(dataset=test_dataset,\n",
    "                              batch_size=BATCH_SIZE,\n",
    "                              num_workers=1,\n",
    "                              shuffle=True)\n",
    "\n",
    "print('expected iterations', math.ceil(len(test_dataset)/BATCH_SIZE))\n",
    "\n",
    "iter_obj = iter(train_dataloader)\n",
    "idx = 0\n",
    "while True:\n",
    "    print(next(iter_obj)[0].shape, idx, '\\r', end='')\n",
    "    idx+=1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2324209",
   "metadata": {},
   "source": [
    "## Infinite Dataloader\n",
    "Create dataloader that runs infinitely by capturing the stop iteration error and intitating a fresh iterator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6e6e795a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# extends pytorch dataloader with infinite iteration\n",
    "class InfiniteDataLoader(DataLoader):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        # Initialize an iterator over the dataset.\n",
    "        self.dataset_iterator = super().__iter__()\n",
    "\n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def __next__(self):\n",
    "        'Expects StopIteration and initiates new iterator'\n",
    "        try:\n",
    "            batch = next(self.dataset_iterator)\n",
    "        except StopIteration:           \n",
    "            self.dataset_iterator = super().__iter__()\n",
    "            batch = next(self.dataset_iterator)\n",
    "        return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f024d292",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset length 3000\n",
      "expected iterations from nomral loader 375\n",
      "torch.Size([8, 3, 224, 224]) 1387 \r"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[7], line 17\u001b[0m\n\u001b[1;32m     15\u001b[0m idx \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m     16\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m---> 17\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;28;43mnext\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43miter_obj\u001b[49m\u001b[43m)\u001b[49m[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mshape, idx, \u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\r\u001b[39;00m\u001b[38;5;124m'\u001b[39m, end\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     18\u001b[0m     idx\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m\n",
      "Cell \u001b[0;32mIn[6], line 14\u001b[0m, in \u001b[0;36mInfiniteDataLoader.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mExpects StopIteration and initiates new iterator\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 14\u001b[0m     batch \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mnext\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdataset_iterator\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m:           \n\u001b[1;32m     16\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset_iterator \u001b[38;5;241m=\u001b[39m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__iter__\u001b[39m()\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/site-packages/torch/utils/data/dataloader.py:628\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    625\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    626\u001b[0m     \u001b[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[39;00m\n\u001b[1;32m    627\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()  \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[0;32m--> 628\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_next_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    629\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    630\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    631\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    632\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called:\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1316\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1313\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_process_data(data)\n\u001b[1;32m   1315\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_shutdown \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_tasks_outstanding \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[0;32m-> 1316\u001b[0m idx, data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_get_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1317\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_tasks_outstanding \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m   1318\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable:\n\u001b[1;32m   1319\u001b[0m     \u001b[38;5;66;03m# Check for _IterableDatasetStopIteration\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1282\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._get_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1278\u001b[0m     \u001b[38;5;66;03m# In this case, `self._data_queue` is a `queue.Queue`,. But we don't\u001b[39;00m\n\u001b[1;32m   1279\u001b[0m     \u001b[38;5;66;03m# need to call `.task_done()` because we don't use `.join()`.\u001b[39;00m\n\u001b[1;32m   1280\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1281\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m-> 1282\u001b[0m         success, data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_try_get_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1283\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m success:\n\u001b[1;32m   1284\u001b[0m             \u001b[38;5;28;01mreturn\u001b[39;00m data\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/site-packages/torch/utils/data/dataloader.py:1120\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_try_get_data\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout\u001b[38;5;241m=\u001b[39m_utils\u001b[38;5;241m.\u001b[39mMP_STATUS_CHECK_INTERVAL):\n\u001b[1;32m   1108\u001b[0m     \u001b[38;5;66;03m# Tries to fetch data from `self._data_queue` once for a given timeout.\u001b[39;00m\n\u001b[1;32m   1109\u001b[0m     \u001b[38;5;66;03m# This can also be used as inner loop of fetching without timeout, with\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1117\u001b[0m     \u001b[38;5;66;03m# Returns a 2-tuple:\u001b[39;00m\n\u001b[1;32m   1118\u001b[0m     \u001b[38;5;66;03m#   (bool: whether successfully get data, any: data if successful else None)\u001b[39;00m\n\u001b[1;32m   1119\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1120\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_data_queue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1121\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m (\u001b[38;5;28;01mTrue\u001b[39;00m, data)\n\u001b[1;32m   1122\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m   1123\u001b[0m         \u001b[38;5;66;03m# At timeout and error, we manually check whether any worker has\u001b[39;00m\n\u001b[1;32m   1124\u001b[0m         \u001b[38;5;66;03m# failed. Note that this is the only mechanism for Windows to detect\u001b[39;00m\n\u001b[1;32m   1125\u001b[0m         \u001b[38;5;66;03m# worker failures.\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/multiprocessing/queues.py:122\u001b[0m, in \u001b[0;36mQueue.get\u001b[0;34m(self, block, timeout)\u001b[0m\n\u001b[1;32m    120\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_rlock\u001b[38;5;241m.\u001b[39mrelease()\n\u001b[1;32m    121\u001b[0m \u001b[38;5;66;03m# unserialize the data after having released the lock\u001b[39;00m\n\u001b[0;32m--> 122\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_ForkingPickler\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mloads\u001b[49m\u001b[43m(\u001b[49m\u001b[43mres\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/site-packages/torch/multiprocessing/reductions.py:305\u001b[0m, in \u001b[0;36mrebuild_storage_fd\u001b[0;34m(cls, df, size)\u001b[0m\n\u001b[1;32m    304\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrebuild_storage_fd\u001b[39m(\u001b[38;5;28mcls\u001b[39m, df, size):\n\u001b[0;32m--> 305\u001b[0m     fd \u001b[38;5;241m=\u001b[39m \u001b[43mdf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdetach\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    306\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    307\u001b[0m         storage \u001b[38;5;241m=\u001b[39m storage_from_cache(\u001b[38;5;28mcls\u001b[39m, fd_id(fd))\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/multiprocessing/resource_sharer.py:57\u001b[0m, in \u001b[0;36mDupFd.detach\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     55\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdetach\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m     56\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m'''Get the fd.  This should only be called once.'''\u001b[39;00m\n\u001b[0;32m---> 57\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[43m_resource_sharer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_connection\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_id\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mas\u001b[39;00m conn:\n\u001b[1;32m     58\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m reduction\u001b[38;5;241m.\u001b[39mrecv_handle(conn)\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/multiprocessing/resource_sharer.py:86\u001b[0m, in \u001b[0;36m_ResourceSharer.get_connection\u001b[0;34m(ident)\u001b[0m\n\u001b[1;32m     84\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mconnection\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Client\n\u001b[1;32m     85\u001b[0m address, key \u001b[38;5;241m=\u001b[39m ident\n\u001b[0;32m---> 86\u001b[0m c \u001b[38;5;241m=\u001b[39m \u001b[43mClient\u001b[49m\u001b[43m(\u001b[49m\u001b[43maddress\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mauthkey\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprocess\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcurrent_process\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mauthkey\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     87\u001b[0m c\u001b[38;5;241m.\u001b[39msend((key, os\u001b[38;5;241m.\u001b[39mgetpid()))\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m c\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/multiprocessing/connection.py:513\u001b[0m, in \u001b[0;36mClient\u001b[0;34m(address, family, authkey)\u001b[0m\n\u001b[1;32m    510\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mauthkey should be a byte string\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m    512\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m authkey \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 513\u001b[0m     \u001b[43manswer_challenge\u001b[49m\u001b[43m(\u001b[49m\u001b[43mc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mauthkey\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    514\u001b[0m     deliver_challenge(c, authkey)\n\u001b[1;32m    516\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m c\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/multiprocessing/connection.py:760\u001b[0m, in \u001b[0;36manswer_challenge\u001b[0;34m(connection, authkey)\u001b[0m\n\u001b[1;32m    758\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m message[:\u001b[38;5;28mlen\u001b[39m(CHALLENGE)] \u001b[38;5;241m==\u001b[39m CHALLENGE, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mmessage = \u001b[39m\u001b[38;5;132;01m%r\u001b[39;00m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m%\u001b[39m message\n\u001b[1;32m    759\u001b[0m message \u001b[38;5;241m=\u001b[39m message[\u001b[38;5;28mlen\u001b[39m(CHALLENGE):]\n\u001b[0;32m--> 760\u001b[0m digest \u001b[38;5;241m=\u001b[39m \u001b[43mhmac\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnew\u001b[49m\u001b[43m(\u001b[49m\u001b[43mauthkey\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmessage\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mmd5\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mdigest()\n\u001b[1;32m    761\u001b[0m connection\u001b[38;5;241m.\u001b[39msend_bytes(digest)\n\u001b[1;32m    762\u001b[0m response \u001b[38;5;241m=\u001b[39m connection\u001b[38;5;241m.\u001b[39mrecv_bytes(\u001b[38;5;241m256\u001b[39m)        \u001b[38;5;66;03m# reject large message\u001b[39;00m\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/hmac.py:170\u001b[0m, in \u001b[0;36mnew\u001b[0;34m(key, msg, digestmod)\u001b[0m\n\u001b[1;32m    153\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnew\u001b[39m(key, msg\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, digestmod\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m):\n\u001b[1;32m    154\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Create a new hashing object and return it.\u001b[39;00m\n\u001b[1;32m    155\u001b[0m \n\u001b[1;32m    156\u001b[0m \u001b[38;5;124;03m    key: bytes or buffer, The starting key for the hash.\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    168\u001b[0m \u001b[38;5;124;03m    or hexdigest() methods.\u001b[39;00m\n\u001b[1;32m    169\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 170\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mHMAC\u001b[49m\u001b[43m(\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmsg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdigestmod\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/transformer/lib/python3.9/hmac.py:61\u001b[0m, in \u001b[0;36mHMAC.__init__\u001b[0;34m(self, key, msg, digestmod)\u001b[0m\n\u001b[1;32m     59\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_digest_cons \u001b[38;5;241m=\u001b[39m digestmod\n\u001b[1;32m     60\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(digestmod, \u001b[38;5;28mstr\u001b[39m):\n\u001b[0;32m---> 61\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_digest_cons \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mlambda\u001b[39;00m d\u001b[38;5;241m=\u001b[39m\u001b[38;5;124mb\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m: _hashlib\u001b[38;5;241m.\u001b[39mnew(digestmod, d)\n\u001b[1;32m     62\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     63\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_digest_cons \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mlambda\u001b[39;00m d\u001b[38;5;241m=\u001b[39m\u001b[38;5;124mb\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m: digestmod\u001b[38;5;241m.\u001b[39mnew(d)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# test the infinite dataloader\n",
    "classes_to_idx, partition, labels = dataset_preprocess(dataset_path)\n",
    "test_dataset = MyDataset(partition['test'], labels, classes_to_idx)\n",
    "\n",
    "print('dataset length', len(test_dataset))\n",
    "\n",
    "BATCH_SIZE = 8\n",
    "train_dataloader = InfiniteDataLoader(dataset=test_dataset,\n",
    "                              batch_size=BATCH_SIZE,\n",
    "                              num_workers=1,\n",
    "                              shuffle=True)\n",
    "print('expected iterations from nomral loader', math.ceil(len(test_dataset)/BATCH_SIZE))\n",
    "\n",
    "iter_obj = iter(train_dataloader)\n",
    "idx = 0\n",
    "while True:\n",
    "    print(next(iter_obj)[0].shape, idx, '\\r', end='')\n",
    "    idx+=1"
   ]
  }
 ],
 "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}