examples/readers/pytorch_data_api_tiledb_dense.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "This example notebook shows how we can train an [image/digit classification](https://pytorch.org/tutorials/beginner/nn_tutorial.html?highlight=mnist)\n",
    "model based on MNIST dataset, by employing TileDB as a storage engine for our training data and labels. We will first download the MNIST\n",
    "dataset and ingest images and labels in two dense TileDB arrays. Continuing, we will use our TileDB support for PyTorch Dataloader API\n",
    "in order to train a image classifier. First, let's import what we need and download our data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import idx2numpy\n",
    "import numpy as np\n",
    "import tiledb\n",
    "import torch\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Download MNIST dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0.3%"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ../data/MNIST/raw/train-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100.0%\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ../data/MNIST/raw/train-images-idx3-ubyte.gz to ../data/MNIST/raw\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100.0%"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ../data/MNIST/raw/train-labels-idx1-ubyte.gz\n",
      "Extracting ../data/MNIST/raw/train-labels-idx1-ubyte.gz to ../data/MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "19.9%"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ../data/MNIST/raw/t10k-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100.0%\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ../data/MNIST/raw/t10k-images-idx3-ubyte.gz to ../data/MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ../data/MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100.0%"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ../data/MNIST/raw/t10k-labels-idx1-ubyte.gz to ../data/MNIST/raw\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "data_home = os.path.join(os.path.pardir, \"data\")\n",
    "data = torchvision.datasets.MNIST(root=data_home, train=False, download=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Images: (60000, 28, 28) uint8\n",
      "Labels: (60000,) uint8\n"
     ]
    }
   ],
   "source": [
    "images = idx2numpy.convert_from_file(os.path.join(data_home, 'MNIST/raw/train-images-idx3-ubyte'))\n",
    "labels = idx2numpy.convert_from_file(os.path.join(data_home, 'MNIST/raw/train-labels-idx1-ubyte'))\n",
    "\n",
    "print(\"Images:\", images.shape, images.dtype)\n",
    "print(\"Labels:\", labels.shape, labels.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Then we proceed with ingesting images and labels into dense TileDB arrays. Here, we should point out that besides the\n",
    "flexibility of TileDB in defining a schema, i.e., multiple dimensions, multiple attributes, compression etc,\n",
    "we choose to define a simple schema. So, for a numpy array of D number of dimensions we create a dense TileDB array,\n",
    "with the same number of dimensions, and a single attribute of data type numpy float32. Moreover, the\n",
    "tile extend of the 1st dimension should always be equal with the batch size, in order to achieve optimal reads while\n",
    "training. Let's define an ingestion function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def ingest_in_tiledb(data: np.array, batch_size: int, uri: str):\n",
    "    # Equal number of dimensions with the numpy array.\n",
    "    dims = [\n",
    "        tiledb.Dim(\n",
    "            name=\"dim_\" + str(dim),\n",
    "            domain=(0, data.shape[dim] - 1),\n",
    "            tile=data.shape[dim] if dim > 0 else batch_size,\n",
    "            dtype=np.int32,\n",
    "        )\n",
    "        for dim in range(data.ndim)\n",
    "    ]\n",
    "\n",
    "    # TileDB schema\n",
    "    schema = tiledb.ArraySchema(\n",
    "        domain=tiledb.Domain(*dims),\n",
    "        sparse=False,\n",
    "        attrs=[tiledb.Attr(name=\"features\", dtype=data.dtype)],\n",
    "    )\n",
    "    # Create array\n",
    "    tiledb.Array.create(uri, schema)\n",
    "\n",
    "    # Ingest\n",
    "    with tiledb.open(uri, \"w\") as tiledb_array:\n",
    "        tiledb_array[:] = {\"features\": data}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we proceed with ingestion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/george/PycharmProjects/TileDB-ML/.venv/lib/python3.9/site-packages/tiledb/ctx.py:448: UserWarning: tiledb.default_ctx and scope_ctx will not function correctly due to bug in IPython contextvar support.  You must supply a Ctx object to each function for custom configuration options. Please consider upgrading to ipykernel >= 6!Please see https://github.com/TileDB-Inc/TileDB-Py/issues/667 for more information.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "data_dir = os.path.join(data_home, 'readers', 'pytorch', 'dense')\n",
    "os.makedirs(data_dir, exist_ok=True)\n",
    "\n",
    "# Ingest images\n",
    "training_images = os.path.join(data_dir, 'training_images')\n",
    "if not os.path.exists(training_images):\n",
    "    ingest_in_tiledb(data=images, batch_size=64, uri=training_images)\n",
    "\n",
    "# Ingest labels\n",
    "training_labels = os.path.join(data_dir, 'training_labels')\n",
    "if not os.path.exists(training_labels):\n",
    "    ingest_in_tiledb(data=labels, batch_size=64, uri=training_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can now explore our TileDB arrays and check their structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ArraySchema(\n",
      "  domain=Domain(*[\n",
      "    Dim(name='dim_0', domain=(0, 59999), tile=64, dtype='int32'),\n",
      "    Dim(name='dim_1', domain=(0, 27), tile=28, dtype='int32'),\n",
      "    Dim(name='dim_2', domain=(0, 27), tile=28, dtype='int32'),\n",
      "  ]),\n",
      "  attrs=[\n",
      "    Attr(name='features', dtype='uint8', var=False, nullable=False),\n",
      "  ],\n",
      "  cell_order='row-major',\n",
      "  tile_order='row-major',\n",
      "  capacity=10000,\n",
      "  sparse=False,\n",
      ")\n",
      "\n",
      "ArraySchema(\n",
      "  domain=Domain(*[\n",
      "    Dim(name='dim_0', domain=(0, 59999), tile=64, dtype='int32'),\n",
      "  ]),\n",
      "  attrs=[\n",
      "    Attr(name='features', dtype='uint8', var=False, nullable=False),\n",
      "  ],\n",
      "  cell_order='row-major',\n",
      "  tile_order='row-major',\n",
      "  capacity=10000,\n",
      "  sparse=False,\n",
      ")\n",
      "\n"
     ]
    }
   ],
   "source": [
    "images_array = tiledb.open(training_images)\n",
    "labels_array = tiledb.open(training_labels)\n",
    "\n",
    "print(images_array.schema)\n",
    "print(labels_array.schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can easily now slice our data and create some plots. We can either slice an image or a part of\n",
    "an image. Because we use only one attribute, we always slice with attribute with index equal to 0.\n",
    "Some examples below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<matplotlib.image.AxesImage at 0x123788ca0>"
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": "<Figure size 432x288 with 2 Axes>",
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXAAAAC4CAYAAAD61bdSAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8qNh9FAAAACXBIWXMAAAsTAAALEwEAmpwYAAAVG0lEQVR4nO3de6xV5ZnH8e9PlDGi8VJHdPBCdSiGGsWWYqcSi1G8xUvxFsnEy4x61EgqiWXiqKm2o60TrXUQp0qVAsZSO6MoUaZCvJSatoxIsXLRisbqoVSsiOAtFnjmj72Oczx7bfZ7ztm3tc/vk5yctZ797L3exVnn4T3vuryKCMzMrHh2aHYDzMysb1zAzcwKygXczKygXMDNzArKBdzMrKBcwM3MCqpfBVzSSZJelrRG0jW1apSZmVWnvl4HLmkQ8AdgAtAJPAdMiohV23mPLzq3uooINWO7PrZrY9CgQU3Z7j777JOUt8MO6X3enXfeOSnvjTfeqJqzZcsWtm3bVnZs75jcmnJjgTUR8RqApJ8BZwAVC7iZDTy9Kcq77757HVtSWUdHR1LeLrvskvyZhx56aFLeFVdcUTXn7bffzo33ZwhlGPBmt/XOLPYZkjokLZW0tB/bMquJasN+kv5G0oPZ60skDW9CM82S1P0kZkTMiIgxETGm3tsy255s2O8u4GRgFDBJ0qgeaRcD70bE3wM/BP69sa00S9efAr4WOKDb+v5ZzKxVfTrsFxGfAF3Dft2dAczOlv8bOE5SU8bVzarpTwF/Dhgh6fOSBgPnAfNr0yyzukgZ9vs0JyK2AO8Bn8v7MA8PWrP1+SRmRGyRNBl4AhgEzIyIlTVrmVmLi4gZwAzwVSjWHP25CoWIWAAsqFFbzOotZdivK6dT0o7A7sA7jWmeWe/4TkwbSFKG/eYDF2bLZwNPhR+aby2qzzfy9Glj/jPT6qzajTySTgHu4P+H/W6W9F1gaUTMl7QzcD9wJLABOK/rXocqn9u0Y/vAAw9Myhs8eHBS3tFHH52UN27cuKS8PffcMykP4KyzzkrObXWdnZ1Jec8991zVnKlTp7JmzZqa3shjVjh5w34R8e1uyx8D5zS6XWZ94SEUM7OCcgE3MysoF3Azs4JyATczKygXcDOzgnIBNzMrKBdwM7OCcgE3Myso38hj1oJGjx6dnPvMM88k5TVrtpt2sm3btuTc6667Lilv8+bNVXPeeSf/cTzugZuZFZQLuJlZQbmAm5kVlAu4mVlBuYCbmRWUC7iZWUH16zJCSa8Dm4GtwJaIGFOLRrW7QYMGlcX6e4nX5MmTc+O77LJLWWzkyJG5uVdeeWVZ7LbbbsvNnTRpUlns448/zs295ZZbymLf+c53cnPNLF0teuDHRsRoF29rdZIOkPS0pFWSVkq6KidnvKT3JC3Pvr6d91lmrcA38thAsgW4OiKWSdoNeF7SoohY1SPvVxFxahPaZ9Yr/e2BB7BQ0vOSOmrRILN6iYh1EbEsW94MrAaGNbdVZn3X3x74uIhYK2kfYJGklyJicfeErLC7uFtLkTSc0sTFS3Je/gdJLwB/Ar4VESsrfEbdju033ngjObfSbdY9tdOt9EuW5P3Yym3cuDEp79hjj03K++STT5LyAObMmZOc21f96oFHxNrs+3pgHjA2J2dGRIzxGLm1Ckm7Ag8BUyJiU4+XlwEHRcQRwJ3AI5U+x8e2NVufe+CShgA7RMTmbPkE4Ls1a1kLOPDAA8tigwcPzs392te+VhYbN25cbu4ee+xRFjvrrLN617h+6OzszI1PmzatLDZx4sTc3LwH8Lzwwgu5ub/85S970br6krQTpeL9QEQ83PP17gU9IhZI+k9Je0fEXxrZTrMU/RlCGQrMk9T1OT+NiF/UpFVmdaDSwXofsDoibq+Qsy/wVkSEpLGU/kpNG6Mwa7A+F/CIeA04ooZtMau3o4HzgRclLc9i1wIHAkTE3cDZwBWStgAfAedFRDShrWZV+TJCGzAi4llAVXKmA9Mb0yKz/vGt9GZmBeUeOJVnP3nqqafKYkW7FCtvBpHrr78+N/f9998viz3wwAO5uevWrSuLvfvuu7m5L7/88vaaaGZ95B64mVlBuYCbmRWUh1DMWtCGDRuSc6+++uqkvNNOOy0p73e/+11S3p133pmU1xvLly9PyjvuuOOS8j744IOkvC9+8YtJeVOmTEnKaxT3wM3MCsoF3MysoDyEQuUHB+U9JKiRV6FUemBP3gN6Kj2MJ+/hO/fff3+/2mVmrcE9cDOzgnIBNzMrKBdwM7OCcgE3Myson8Sk8jW3U6dOLYudemr+VIl5187mPV+7krzrXydMmJCbm3dta6XrWK+6qmzeXjNrE+6Bm5kVlHvgZgX3yCOPJOXlPZwtT95sS3mOOCJtOoBLLrkkKQ/gtttuS8pLvcMy1cqVudOelrn00ktrut3+cg/czKygXMBtwJH0uqQXJS2XtDTndUmaJmmNpN9L+lIz2mlWjYdQbKA6djsTFZ8MjMi+jgJ+lH03aylVC7ikmcCpwPqIOCyL7QU8CAwHXgfOjYj8p/kXWN7YYqVxxLxxw0pjhBdffHFZLG/srzfjfJXG8Do6OpI/wz51BjAnmwvzt5L2kLRfRJTPYmHWRClDKLOAk3rErgGejIgRwJPZullRBLBQ0vOS8v6HGwa82W29M4t9hqQOSUvzhmHMGqFqAY+IxUDPC6XPAGZny7OBb9S2WWZ1NS4ivkRpqORKScf05UMiYkZEjImIMbVtnlmavp7EHNrtz8k/A0MrJbqXYq0mItZm39cD84CxPVLWAgd0W98/i5m1lH5fhZKNE8Z2XncvxVqGpCGSdutaBk4AVvRImw9ckF2N8lXgPY9/Wyvq61Uob3Wd1JG0H7C+lo1qZZs2bUrOfe+995Jz824QePDBB3Nz82aat2RDgXmSoHT8/zQifiHpcoCIuBtYAJwCrAE+BP6pSW01266+FvD5wIXALdn3R2vWIrM6iojXgLLLg7LC3bUcwJWNbFcj9KbzkaI3HZRUqXc6zp07Nymv3Ts7VYdQJM0FfgOMlNQp6WJKhXuCpFeA47N1MzNroKo98IiYVOGltGmhzcysLnwrvZlZQbmAm5kVlJ+FUkc33nhjbvzLX/5yWezrX/96Wez444/Pff/ChQv71S4zaw/ugZuZFZQLuJlZQbmAm5kVlAu4mVlBqXTTWYM2JjVuYy3skEMOKYstW7asLLZx48bc9z/99NNlsaVL858Vdtddd5XFGvkzb7SIUDO2OxCP7SFDhiTlPf7448mfmXcyP8+JJ56YlNdOJ/zzjm33wM3MCsoF3MysoFzAzcwKygXczKygfCdmE7z66qtlsYsuuqgs9pOf/CT3/eeff35SDPJPNM2ZMyc3d906z1lgViTugZuZFZQLuJlZQbmA24AhaaSk5d2+Nkma0iNnvKT3uuV8u0nNNavKY+A2YETEy8BoAEmDKM00Py8n9VcRcWoDm2bWJ+6B20B1HPBqRPyx2Q0x66uqt9JLmgmcCqyPiMOy2I3ApcDbWdq1EbGg6sYG4O3G/XHYYYflxm+//fay2HHHpc9wd8899+TGb7755rLY2rVrkz+3FaTeSp8d18siYnqP+HjgIaAT+BPwrYhYWeEzOoCObLX8Ie8G5D86opLly5cn5VV6zERPTz31VFJepUdR9DR9+vTqSZlaP7Kir7fSzwJOyon/MCJGZ19Vi7dZq5A0GDgd+K+cl5cBB0XEEcCdwCOVPiciZkTEmIgYU5eGmlVRtYBHxGJgQwPaYtYoJ1Pqfb/V84WI2BQR72fLC4CdJO3d6AaapejPGPhkSb+XNFPSnpWSJHVIWiop7W8Us/qbBMzNe0HSvpKULY+l9DvyTgPbZpasrwX8R8AhlM7orwN+UCnRf2ZaK5E0BJgAPNwtdrmky7PVs4EVkl4ApgHnRTs/f9cKrU+XEXb/01PSj4HHatYi+9SKFSty4+eee25Z7LTTTsvNzbsd/7LLLsvNHTFiRFlswoQJ22ti4UTEB8DnesTu7rY8HUg/U2XWRH3qgUvar9vqRCC/0piZWd1U7YFLmguMB/aW1AncAIyXNBoI4HUgv0tnZmZ1U7WAR8SknPB9dWiLmZn1gu/ENDMrKD8LxczqKu/595VccMEFSXmzZ8+u6eel5qVO5AzpbezPc/hdwAso7zbi+++/Pzf33nvvLYvtuGP+j/2YY44pi40fPz4395lnnqnYPjNrDA+hmJkVlAu4mVlBuYCbmRWUC7iZWUH5JGYLO/zww3PjZ599dlnsK1/5Sm5upROWeVatWlUWW7x4cfL7zayx3AM3MysoF3Azs4JyATczKyiPgZtZy5g3b15S3iuvvJKUd8cddyTlpc4p+/3vfz8pD+Cggw5Kyrvpppuq5qxfvz437h64mVlBuQfeBCNHjiyLTZ48uSx25pln5r5/33337df2t27dmhvPeybDtm3b+rWtZslmnT8VWB8Rh2WxvYAHgeGUHoN8bkS8m/PeC4Hrs9WbIiLtoRZmDeYeuLWrWcBJPWLXAE9GxAjgyWz9M7IifwNwFDAWuGF7c76aNZMLuLWliFgMbOgRPgPo6k3PBr6R89YTgUURsSHrnS+i/D8Cs5bgIRQbSIZGRNc40Z+BoTk5w4A3u613ZrEykjqAjpq20KwXXMBtQIqIkNSv2eYjYgYwA6C/n2XWFylzYh4AzKHUWwlgRkT8R+oJoYEi78TipEl5s9Hln7AcPnx4rZsEwNKlS8tiN998c27u/Pnz69KGFvKWpP0iYl02MXfetVlrKc0B22V/4JkGtM2s11LGwLcAV0fEKOCrwJWSRpFwQsisxcwHLsyWLwQezcl5AjhB0p7ZycsTsphZy6lawCNiXUQsy5Y3A6spjQmmnBAyawpJc4HfACMldUq6GLgFmCDpFeD4bB1JYyTdCxARG4B/A57Lvr6bxcxaTq/GwCUNB44ElpB2QsgneqwpIiJ//ArKbrmLiKXAJd3WZwIz69Q0q4EVK1Yk5eU9uTPP6aefnpQ3a9aspDyAyy+/PCnvC1/4QtWcK664IjeefBmhpF2Bh4ApEbGp+2sREZTGx8tExIyIGBMRY1K3ZWZm1SUVcEk7USreD0TEw1n4rexEENs5IWRmZnWSchWKgPuA1RFxe7eXuk4I3ULlE0KFNnRo+ajQqFGjcnOnT59eFjv00ENr3iaAJUuWlMVuvfXW3NxHHy3/sRT19ngz+6yUMfCjgfOBFyUtz2LXUircP89ODv0ROLcuLTQzs1xVC3hEPAuowstpz2A0M7Oa87NQzMwKygXczKygBtyzUPbaa6+y2D333JObO3r06LLYwQcfXOsmAfDrX/+6LPaDH/wgN/eJJ8pvDPzoo49q3iYza23ugZuZFdSA64Gb2cCxcePGpLw5c+Yk5d13333J295xx7Tyeswxx1TN2XXXXXPj7oGbmRWUC7iZWUG5gJuZFVRbjIEfddRRZbGpU6fm5o4dO7YsNmxY7oxZ/fbhhx/mxqdNm1YW+973vlcW++CDD2reJjNrH+6Bm5kVlAu4mVlBuYBb25E0U9J6SSu6xW6V9JKk30uaJ2mPCu99XdKLkpZLKp9Q1KyFuIBbO5oFnNQjtgg4LCIOB/4A/Ot23n9sRIz2JCTW6triJObEiROTYr21atWqsthjjz2Wm7tly5ayWKVb4VNvLrC+iYjF2fR/3WMLu63+Fkiba8ushbkHbgPRPwP/U+G1ABZKej6bz9WsZbVFD9wslaTrgC3AAxVSxkXEWkn7AIskvRQRiyt8lifsbpLDDz88Ke+cc85Jysu7vDhP6u3xvZH3l35PH3/8cW7cPXAbMCRdBJwK/GM2EXeZiFibfV8PzAMq/mZ7wm5rNhdwGxAknQT8C3B6ROTeYSVpiKTdupaBE4AVeblmraBqAZd0gKSnJa2StFLSVVn8Rklrs8utlks6pf7NNatO0lzgN8BISZ3ZvK3Tgd0oDYssl3R3lvt3khZkbx0KPCvpBeB/gccj4hdN2AWzJCkDOluAqyNiWdY7eV7Souy1H0bEbfVrXpprrrkmKWYDQ0RMygnnPgc0Iv4EnJItvwYcUcemmdVUyqTG64B12fJmSauB+jw8xMzMkvVqDDy7tvZIYEkWmpzd2TZT0p4V3tMhaanvajMzq63kAi5pV+AhYEpEbAJ+BBwCjKbUQ8+9a8Vn6s3M6iOpgEvaiVLxfiAiHgaIiLciYmtEbAN+zHYutzIzs9pLuQpFlE4ArY6I27vF9+uWNhFfbmVm1lApV6EcDZwPvChpeRa7FpgkaTSlW49fBy6rQ/vMrOBGjhyZnPvNb34zKe/MM89Mytt3332Tt11rW7duTcpbt25d1Zy//vWvufGUq1CeBZTz0oKcmJmZNYjvxDQzKygXcDOzgnIBNzMrKBdwM7OCcgE3MysoF3Azs4JyATczKygXcDOzglKFmaXqszHpbeCP2erewF8atvHG8X41z0ER8bfN2HCPY7tLEf7NUrTLfkBx9yX32G5oAf/MhqWl7fiEQu+XdWmXf7N22Q9or30BD6GYmRWWC7iZWUE1s4DPaOK268n7ZV3a5d+sXfYD2mtfmjcGbmZm/eMhFDOzgnIBNzMrqIYXcEknSXpZ0hpJ1zR6+7Ukaaak9ZJWdIvtJWmRpFey73s2s419IekASU9LWiVppaSrsnjh960R2uwYf13Si5KWS1ra7Pb0Rrv+fnbX0AIuaRBwF3AyMIrStGyjGtmGGpsFnNQjdg3wZESMAJ7M1otmC3B1RIwCvgpcmf2c2mHf6qoNj3GAYyNidAGvn55Fe/5+fqrRPfCxwJqIeC0iPgF+BpzR4DbUTEQsBjb0CJ8BzM6WZwPfaGSbaiEi1kXEsmx5M7AaGEYb7FsDtNUxXmTt+vvZXaML+DDgzW7rnVmsnQyNiK5ZSv8MDG1mY/pL0nDgSGAJbbZvddJux3gACyU9L6mj2Y2pgbY6hlNmpbc+ioiQVNjrNCXtCjwETImITdL/z21d9H2zZOMiYq2kfYBFkl7KeraF1w7HcKN74GuBA7qt75/F2slbkvYDyL6vb3J7+kTSTpSK9wMR8XAWbot9q7O2OsYjYm32fT0wj9IQUZG11THc6AL+HDBC0uclDQbOA+Y3uA31Nh+4MFu+EHi0iW3pE5W62vcBqyPi9m4vFX7fGqBtjnFJQyTt1rUMnACs2P67Wl5bHcMNvxNT0inAHcAgYGZE3NzQBtSQpLnAeEqPqHwLuAF4BPg5cCClx4ueGxE9T6S0NEnjgF8BLwLbsvC1lMbBC71vjdAux7ikgyn1uqE03PrTIu1Lu/5+dudb6c3MCsp3YpqZFZQLuJlZQbmAm5kVlAu4mVlBuYCbmRWUC7iZWUG5gJuZFdT/AS9b150Xyes9AAAAAElFTkSuQmCC\n"
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot an image\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(images_array[0][images_array.schema.attr(0).name], cmap=\"gray\")\n",
    "\n",
    "# Plot part of the same image\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(images_array[0, 5:20, 5:20][images_array.schema.attr(0).name], cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's move on and create a TileDB loader that extends the PyTorch DataLoader API for training a machine learning model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 1 Batch: 0 Loss: 2.304748\n",
      "Train Epoch: 1 Batch: 100 Loss: 2.277155\n",
      "Train Epoch: 1 Batch: 200 Loss: 2.203359\n",
      "Train Epoch: 1 Batch: 300 Loss: 1.895098\n",
      "Train Epoch: 1 Batch: 400 Loss: 1.497304\n",
      "Train Epoch: 2 Batch: 0 Loss: 1.435658\n",
      "Train Epoch: 2 Batch: 100 Loss: 1.305221\n",
      "Train Epoch: 2 Batch: 200 Loss: 0.990590\n",
      "Train Epoch: 2 Batch: 300 Loss: 1.103210\n",
      "Train Epoch: 2 Batch: 400 Loss: 0.903957\n"
     ]
    }
   ],
   "source": [
    "import multiprocessing\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "from tiledb.ml.readers.pytorch import PyTorchTileDBDataLoader\n",
    "from tiledb.ml.readers.types import ArrayParams\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self, shape):\n",
    "        super(Net, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(np.product(shape), 32),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(32, 32),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(32, 10),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "def do_random_noise(img, mag=0.1):\n",
    "    noise = np.random.uniform(-1, 1,img.shape)*mag\n",
    "    img = img + noise\n",
    "    img = np.clip(img,0,1)\n",
    "    return img\n",
    "\n",
    "ctx = tiledb.Ctx({'sm.mem.total_budget': 1024**2})\n",
    "with tiledb.open(training_images, ctx=ctx) as x, tiledb.open(training_labels, ctx=ctx) as y:\n",
    "    # Because of this issue (https://github.com/pytorch/pytorch/issues/59451#issuecomment-854883855) we avoid using multiple workers on Jupyter.\n",
    "    train_loader = PyTorchTileDBDataLoader(\n",
    "        ArrayParams(x, fn=do_random_noise),\n",
    "        ArrayParams(y),\n",
    "        batch_size=128,\n",
    "        num_workers=0,\n",
    "        shuffle_buffer_size=256,\n",
    "    )\n",
    "\n",
    "    net = Net(shape=(28, 28))\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.5)\n",
    "\n",
    "    for epoch in range(1, 3):\n",
    "        net.train()\n",
    "        for batch_idx, (inputs, labels) in enumerate(train_loader):\n",
    "            # zero the parameter gradients\n",
    "            optimizer.zero_grad()\n",
    "            # forward + backward + optimize\n",
    "            outputs = net(inputs.to(torch.float))\n",
    "            loss = criterion(outputs, labels.to(torch.float).type(torch.LongTensor))\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            if batch_idx % 100 == 0:\n",
    "                print('Train Epoch: {} Batch: {} Loss: {:.6f}'.format(\n",
    "                epoch, batch_idx, loss.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}