dask_array_dataloader.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f48c0293-fe0a-4687-9a02-7826d52d80a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.dlpack import from_dlpack\n",
    "import queue\n",
    "import threading\n",
    "\n",
    "import dask.array as da\n",
    "import cupy\n",
    "import numpy as np\n",
    "import time\n",
    "import math\n",
    "import os\n",
    "\n",
    "\n",
    "class BatchQueue:\n",
    "    def __init__(self, darr, batch_size=32):\n",
    "        self.darr = darr\n",
    "        self.batch_size = batch_size\n",
    "        self.input_queue = queue.Queue()\n",
    "        self.output_queue = queue.Queue()\n",
    "        self.worker = threading.Thread(\n",
    "            target=self.reader,\n",
    "            daemon=True,\n",
    "        )      \n",
    "\n",
    "    def __len__(self):\n",
    "        # Number of batches in self.darr\n",
    "        return math.ceil(self.darr.shape[0] / self.batch_size)   \n",
    "\n",
    "    def start_worker(self):\n",
    "        self.worker.start()\n",
    "\n",
    "    def put(self, batch_index):\n",
    "        # Add a new batch index to input_queue\n",
    "        start_index = batch_index * self.batch_size\n",
    "        self.input_queue.put(start_index)\n",
    "\n",
    "    def get(self):\n",
    "        # Get a pytorch tensor from output_queue\n",
    "        return self.output_queue.get()\n",
    "\n",
    "    def reader(self):\n",
    "        # Function to be executed by the IO worker\n",
    "        while True:\n",
    "            start_index = self.input_queue.get()\n",
    "            self.output_queue.put(\n",
    "                torch.as_tensor(\n",
    "                    self.darr[\n",
    "                        start_index:start_index + self.batch_size\n",
    "                    ].compute(scheduler=\"synchronous\")\n",
    "                )\n",
    "            )\n",
    "\n",
    "        \n",
    "class DaskDataset(torch.utils.data.IterableDataset):\n",
    "    def __init__(self, darr, batch_size=32, prefetch=True):\n",
    "        super(DaskDataset).__init__()\n",
    "        self.darr = darr\n",
    "        self.batch_size = batch_size\n",
    "        self.prefetch = prefetch\n",
    "        self.batch_queue = BatchQueue(self.darr, batch_size=batch_size)\n",
    "\n",
    "    def __len__(self):\n",
    "        # Total legth of 0th index\n",
    "        return self.darr.shape[0]\n",
    "\n",
    "    def __iter__(self):\n",
    "        # Simple iteration over batches\n",
    "        if self.prefetch:\n",
    "            # Let BatchQueue do all the work\n",
    "            nbatches = len(self.batch_queue)\n",
    "            if nbatches:\n",
    "                self.batch_queue.start_worker()\n",
    "                self.batch_queue.put(0)\n",
    "            for batch_id in range(0, nbatches):\n",
    "                if batch_id < nbatches - 1:\n",
    "                    # Pre-fetch the next batch\n",
    "                    self.batch_queue.put(batch_id + 1)\n",
    "                yield self.batch_queue.get()    \n",
    "        else:\n",
    "            # Iterate over batches directly\n",
    "            for start in range(0, len(self), self.batch_size):\n",
    "                yield torch.as_tensor(\n",
    "                    self.darr[start:start+self.batch_size].compute(scheduler=\"synchronous\")\n",
    "                )    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1c488db0-9816-44f0-9158-c63dd51f1224",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate chunked dask arrays of mamy cupy random arrays\n",
    "total_dataset_size = 1000\n",
    "batch_size = 32\n",
    "\n",
    "path = \"./zarr_data\"\n",
    "if not os.path.isdir(path):\n",
    "    rs = da.random.RandomState(RandomState=np.random.RandomState)\n",
    "    x = rs.normal(10, 1, size=(total_dataset_size, 256, 256, 3), chunks=(batch_size, 256, 256, 3))\n",
    "    x.to_zarr(path)\n",
    "arr = da.from_zarr(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "19d585f0-a6f2-4479-8b7f-e4e25b346aa8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.041147708892822266\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025266408920288086\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025189638137817383\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025146484375\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.026011228561401367\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.02548503875732422\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025165319442749023\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025147199630737305\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025369644165039062\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025413990020751953\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025143861770629883\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025151491165161133\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.02539539337158203\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025122880935668945\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025143146514892578\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025171518325805664\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025405406951904297\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.02515387535095215\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025145292282104492\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025151729583740234\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025396347045898438\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025168180465698242\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025168895721435547\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025147199630737305\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025513648986816406\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025143861770629883\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025127172470092773\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.02512526512145996\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.02537822723388672\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.02514362335205078\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.025150775909423828\n",
      "<class 'torch.Tensor'> torch.Size([8, 256, 256, 3]) 0.025113821029663086\n",
      "CPU times: user 121 ms, sys: 354 ms, total: 475 ms\n",
      "Wall time: 829 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# WITH Pre-Fetching: Run time <1s\n",
    "tlast = time.time()\n",
    "for batch in DaskDataset(arr, batch_size=batch_size, prefetch=True):\n",
    "    time.sleep(0.025) # Fake \"compute\"\n",
    "    step_time = time.time() - tlast\n",
    "    print(type(batch), batch.size(), step_time)\n",
    "    tlast = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8ef5799b-ba6e-4dc2-8f01-286bebc7db7a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.037016868591308594\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.040343284606933594\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.041562557220458984\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04128861427307129\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04243183135986328\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04186582565307617\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04220390319824219\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04277467727661133\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.044110774993896484\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04320359230041504\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04401397705078125\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.044938087463378906\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04629015922546387\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04532670974731445\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.046631574630737305\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.0467219352722168\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04908108711242676\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.047563791275024414\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.0480196475982666\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04821300506591797\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.0476534366607666\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.046339988708496094\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.047934532165527344\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.0477907657623291\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04935812950134277\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04751396179199219\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04677867889404297\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04653525352478027\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.04577279090881348\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.045320749282836914\n",
      "<class 'torch.Tensor'> torch.Size([32, 256, 256, 3]) 0.045760393142700195\n",
      "<class 'torch.Tensor'> torch.Size([8, 256, 256, 3]) 0.034364938735961914\n",
      "CPU times: user 188 ms, sys: 466 ms, total: 655 ms\n",
      "Wall time: 1.44 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# WITHOUT Pre-Fetching: Run time >1s\n",
    "tlast = time.time()\n",
    "for batch in DaskDataset(arr, batch_size=batch_size, prefetch=False):\n",
    "    time.sleep(0.025) # Fake \"compute\"\n",
    "    step_time = time.time() - tlast\n",
    "    print(type(batch), batch.size(), step_time)\n",
    "    tlast = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38daf91e-40b1-4c2c-8b1e-68efe67cc065",
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}