examples/readers/tensorflow_data_api_tiledb_dense.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example notebook shows how we can train an image classification model, as described [here](https://github.com/tensorflow/docs/blob/master/site/en/tutorials/quickstart/beginner.ipynb),\n",
    "with the use of TileDB support for Tensorflow Data API for dense TileDB arrays. We will firstly ingest our MNIST dataset in two dense TileDB arrays, i.e, x and y,\n",
    "and then move on with training of a classification model with Keras. Firstly, let's import what we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "import tiledb\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load MNIST dataset for Keras datasets and scale."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Images: (60000, 28, 28) float64\n",
      "Labels: (60000,) uint8\n"
     ]
    }
   ],
   "source": [
    "(images, labels), _ = tf.keras.datasets.mnist.load_data()\n",
    "images = images / 255.0\n",
    "\n",
    "print(\"Images:\", images.shape, images.dtype)\n",
    "print(\"Labels:\", labels.shape, labels.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": 3,
   "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": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now we proceed with ingestion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "data_dir = os.path.join(os.path.pardir, 'data', 'readers', 'tensorflow', '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": {},
   "source": [
    "We can now explore our TileDB arrays and check their structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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='float64', 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": {},
   "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": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f2b50eab990>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXAAAAC4CAYAAAD61bdSAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8qNh9FAAAACXBIWXMAAAsTAAALEwEAmpwYAAAVB0lEQVR4nO3de4zddZnH8feHQmO4RGBZWgQKQmoRCRSthbUES7BYSAG5hmYD6CIDhEZIsFkWjOJGlISLWMoCFbCUYMUNFBroCoSLlahdSh2gF5BLEIZOKAKl5RYsffaP86s7zvke5jsz5/Y783klkznnOc855zmd3zz9zvd3+SoiMDOz8tmq1QWYmdnQuIGbmZWUG7iZWUm5gZuZlZQbuJlZSbmBm5mV1LAauKTpkp6T9IKki+tVlJmZDUxDPQ5c0ijgz8A0oAd4ApgZEas/4Tk+6NwaKiLUivf1tl0fo0aNasn77rrrrll5W22VP+b91Kc+lZX3yiuvDJizadMmNm/eXLVtb51dTbXJwAsR8RKApF8BxwM1G7iZjTyDacqf/vSnG1hJbV1dXVl52267bfZr7rfffll555133oA5b7zxRjI+nCmU3YFX+9zvKWL/QFKXpOWSlg/jvczqYqBpP1XMKR5/WtIXW1GnWY7hNPDUn6pVf0ZGxLyImBQRk4bxXmbDVkz7XQ8cDewPzJS0f7+0o4HxxVcXcENTizQbhOE08B5gzz739wDWDq8cs4b6+7RfRHwEbJn26+t4YEFU/BHYUdJuzS7ULMdwGvgTwHhJn5U0GjgNWFyfsswaImfaL2tqEDw9aK035J2YEbFJ0izgAWAUcGtErKpbZWb1lzPtlzU1CJXpQWAe+CgUa43hHIVCRCwBltSpFrNGy5n289SglYbPxLSRJGfabzFwRnE0yqHAOxHR2+xCzXIMawRuVia1pv0knVs8fiOVvyiPAV4A3ge+1ap6c40bNy4rb/To0Vl5U6ZMyco77LDDsvJ22mmnrDyAk046KTu33fX09GTlzZ07d8Cc2bNnJ+Nu4DaipKb9isa95XYA5ze7LrOh8BSKmVlJuYGbmZWUG7iZWUm5gZuZlZQbuJlZSbmBm5mVlBu4mVlJuYGbmZWUT+Qxa0MTJ07Mzn3sscey8lq12k0n2bx5c3bupZdempW3cePGAXPefPPNZNwjcDOzknIDNzMrKTdwM7OScgM3MyspN3Azs5JyAzczK6lhHUYo6WVgI/AxsCkiJtWjqE43atSoqthwD/GaNWtWMr7ttttWxSZMmJDMPf/86stgX3XVVcncmTNnVsU+/PDDZO4VV1xRFfvhD3+YzDWzfPUYgR8RERPdvK3dSdpT0qOS1khaJemCRM5USe9I6i6+vt+KWs1y+EQeG0k2ARdFxApJOwBPSnooIlb3y/tdRMxoQX1mgzLcEXgAD0p6UlJXPQoya5SI6I2IFcXtjcAaYPfWVmU2dMMdgU+JiLWSdgUekvRsRCztm1A0djd3ayuS9gYOBpYlHv4XSU8Ba4HvRsSqGq/RsG37lVdeyc6tdZp1f510Kv2yZakfW7X169dn5R1xxBFZeR999FFWHsCCBQuyc4dqWCPwiFhbfF8HLAImJ3LmRcQkz5Fbu5C0PXAXcGFEbOj38Apgr4g4CLgOuKfW63jbtlYb8ghc0nbAVhGxsbh9FPCfdausDYwbN64qNnr06GTuV77ylarYYYcdlszdcccdq2InnXTS4Iobhp6enmR8zpw5VbETTjghmZu6AM9TTz2VzP3tb387iOoaS9I2VJr3HRFxd//H+zb0iFgi6b8k7RIRf21mnWY5hjOFMgZYJGnL6/wyIn5Tl6rMGkCVjfUWYE1EXFMjZyzwekSEpMlU/krNm6Mwa7IhN/CIeAk4qI61mDXaFOB04BlJ3UXsEmAcQETcCJwMnCdpE/ABcFpERAtqNRuQDyO0ESMiHgc0QM5cYG5zKjIbHp9Kb2ZWUh6BU3v1k0ceeaQqVrZDsVIriHzve99L5r777rtVsTvuuCOZ29vbWxV7++23k7nPPffcJ5VoZkPkEbiZWUm5gZuZlZSnUMza0FtvvZWde9FFF2XlHXvssVl5f/rTn7Lyrrvuuqy8weju7s7KO/LII7Py3nvvvay8L3zhC1l5F154YVZes3gEbmZWUm7gZmYl5SkUal84KHWRoGYehVLrgj2pC/TUuhhP6uI7t99++7DqMrP24BG4mVlJuYGbmZWUG7iZWUm5gZuZlZR3YlL7mNvZs2dXxWbMSC+VmDp2NnV97VpSx79OmzYtmZs6trXWcawXXFC1bq+ZdQiPwM3MSsojcLOSu+eee7LyUhdnS0mttpRy0EF5ywF8+9vfzsoDuOqqq7Lycs+wzLVqVXLZ0ypnn312Xd93uDwCNzMrKTdwG1EkvSzpGUndkpYnHpekOZJekPS0pC+2ok6zHJ5CsZHoiE9YpPhoYHzxdQhwQ/HdrO0M2MAl3QrMANZFxAFFbGfgTmBv4GXg1IhIX82/xFJzi7XmEVPzhrXmCM8666yqWGrubzDzfLXm8Lq6urJfwwA4HlhQrIP5R0k7StotIqpXsDBrsZwplPnA9H6xi4GHI2I88HBx36wMAnhQ0pOSUv+77Q682ud+TxGrIqlL0vLUVIxZMwzYwCNiKdD/QOnjgduK27cB36hvWWYNMyUivkhlquR8SYf3ezy16HFyVfqImBcRkyJiUr2LNMsx1J2YY7b8SVl837VWokcp1k4iYm3xfR2wCJjcL6UH2LPP/T2Atc2pzmxwGn4Uikcp1i4kbSdphy23gaOAlf3SFgNnFEejHAq84/lva1dDPQrl9S07diTtBqyrZ1HtbMOGDdm577zzTnZu6gSBO++8M5mbWmnesowBFkmCyrb/y4j4jaRzASLiRmAJcAzwAvA+8K0W1Wo2oKE28MXAmcAVxfd761aRWYNExEtA1aFBRePecjuA85tZV7MMZvCRYzADlFy5ZzouXLgwK6/TBzsDTqFIWgj8AZggqUfSWVQa9zRJzwPTivtmZtZEA47AI2JmjYfyloU2M7OG8Kn0ZmYl5QZuZlZSvhZKA1122WXJ+Je+9KWq2Fe/+tWq2Ne+9rXk8x988MFh1WVmncEjcDOzknIDNzMrKTdwM7OScgM3MyspVU48a9KbSc17sza27777VsVWrFhRFVu/fn3y+Y8++mhVbPny9LXCrr/++qpYM3/mzRYRqasJNtxI3La32267rLz7778/+zVTO/NTvv71r2flddIO/9S27RG4mVlJuYGbmZWUG7iZWUm5gZuZlZTPxGyBF198sSr2zW9+syr2i1/8Ivn8008/PSsG6R1NCxYsSOb29nrdArMy8QjczKyk3MDNzErKDdxGDEkTJHX3+dog6cJ+OVMlvdMn5/stKtdsQJ4DtxEjIp4DJgJIGgW8RmVl+v5+FxEzmlia2ZB4BG4j1ZHAixHxl1YXYjZUA55KL+lWYAawLiIOKGKXAWcDbxRpl0TEkgHfbASebjwcBxxwQDJ+zTXXVMWOPDJ/hbubbropGb/88surYq+99lr267aD3FPpi+16RUTM7RefCtwF9ABrge9GxKoar9EFdBV3qy/ybkD60hG1dHd3Z+XVusxEf4888khWXq1LUfQ3d+7cgZMK9b5kxVBPpZ8PTE/EfxoRE4uvAZu3WbuQNBo4DvjvxMMrgL0i4iDgOuCeWq8TEfMiYlJETGpIoWYDGLCBR8RS4K0m1GLWLEdTGX2/3v+BiNgQEe8Wt5cA20japdkFmuUYzhz4LElPS7pV0k61kiR1SVouKe9vFLPGmwksTD0gaawkFbcnU/kdebOJtZllG2oDvwHYl8oe/V7g6lqJ/jPT2omkbYFpwN19YudKOre4ezKwUtJTwBzgtOjk6+9aqQ3pMMK+f3pK+jlwX90qsr9buXJlMn7qqadWxY499thkbup0/HPOOSeZO378+KrYtGnTPqnE0omI94F/6he7sc/tuUD+niqzFhrSCFzSbn3ungCkO42ZmTXMgCNwSQuBqcAuknqAHwBTJU0EAngZSA/pzMysYQZs4BExMxG+pQG1mJnZIPhMTDOzkvK1UMysoVLXv6/ljDPOyMq77bbb6vp6uXm5CzlDfo3DuQ6/G3gJpU4jvv3225O5N998c1Vs663TP/bDDz+8KjZ16tRk7mOPPVazPjNrDk+hmJmVlBu4mVlJuYGbmZWUG7iZWUl5J2YbO/DAA5Pxk08+uSr25S9/OZlba4dlyurVq6tiS5cuzX6+mTWXR+BmZiXlBm5mVlJu4GZmJeU5cDNrG4sWLcrKe/7557Pyrr322qy83DVlf/KTn2TlAey1115ZeT/60Y8GzFm3bl0y7hG4mVlJeQTeAhMmTKiKzZo1qyp24oknJp8/duzYYb3/xx9/nIynrsmwefPmYb1XKxQrzs8A1kXEAUVsZ+BOYG8ql0A+NSLeTjx3OvAzYBRwc0Rc0aSyzQbNI3DrRPOB6f1iFwMPR8R44OHi/j+QNAq4nsqix/sDMyXt39hSzYbODdw6TkQsBd7qFz4e2HJ5uNuAbySeOhl4ISJeioiPgF8VzzNrS27gNlKMiYhegOL7romc3YFX+9zvKWJJkrokLZe0vK6VmmXyHLjZ/1MiVnNF+oiYB8wDkOSV663pctbE3BNYAIwFNgPzIuJnuTuFRorUjsWZM1Or0aV3WO699971LgmA5curB4eXX355Mnfx4sUNqaFNvC5pt4joLRblTh2X1QPs2ef+HsDaplRnNgQ5UyibgIsi4vPAocD5xY6dAXcKmbWRxcCZxe0zgXsTOU8A4yV9VtJo4LTieWZtacAGHhG9EbGiuL0RWENlXjBnp5BZ00laCPwBmCCpR9JZwBXANEnPA9OK+0j6jKQlABGxCZgFPEBlO/91RKxqxWcwyzGoOXBJewMHA8vot1NIUmqnEJK6gK5h1mmWLSLSc1dQdbpdRKwFjulzfwmwpEGlWZ2sXLkyKy915c6U4447Litv/vz5WXkA5557blbe5z73uQFzzjvvvGQ8u4FL2h64C7gwIjZIqf091byjx8ysMbIOI5S0DZXmfUdE3F2EXy92BvEJO4XMzKxBco5CEXALsCYirunz0JadQldQe6dQqY0ZM6Yqtv/+6RPz5s6dWxXbb7/96l4TwLJly6piV155ZTL33nurfyxlPD3ezKrlTKFMAU4HnpHUXcQuodK4f13sIHoFOKUhFZqZWdKADTwiHid9ggMkdgqZmVlz+FR6M7OScgM3MyupEXctlJ133rkqdtNNNyVzJ06cWBXbZ5996l0SAL///e+rYldffXUy94EHHqiKffDBB3Wvyczam0fgZmYlNeJG4GY2cqxfvz4rb8GCBVl5t9xyS/Z7b711Xns9/PDDB8zZfvvtk3GPwM3MSsoN3MyspNzAzcxKqiPmwA855JCq2OzZs5O5kydPrortvnvNVbOG5f3330/G58yZUxX78Y9/XBV777336l6TmXUOj8DNzErKDdzMrKTcwK3jSLpV0jpJK/vErpT0rKSnJS2StGON574s6RlJ3V5t3tqdG7h1ovnA9H6xh4ADIuJA4M/Af3zC84+IiIkRMalB9ZnVRUfsxDzhhBOyYoO1evXqqth9992XzN20aVNVrNap8LknF9jQRMTSYvm/vrEH+9z9I5C31pZZG/MI3EaifwP+p8ZjATwo6cliPVezttURI3CzXJIuBTYBd9RImRIRa4tFuh+S9GxELK3xWl6wu0UOPPDArLxTTslbZyZ1eHFK7unxg5H6S7+/Dz/8MBn3CNxGDElnAjOAf42I5ALbxSr1RMQ6YBFQ8zc7IuZFxCTPlVuruIHbiCBpOvDvwHERkTzDStJ2knbYchs4CliZyjVrBwM2cEl7SnpU0hpJqyRdUMQvk/RacbhVt6RjGl+u2cAkLQT+AEyQ1FOs2zoX2IHKtEi3pBuL3M9IWlI8dQzwuKSngP8F7o+I37TgI5hlyZnQ2QRcFBEritHJk5IeKh77aURc1bjy8lx88cVZMRsZImJmIpy8DmgxZXJMcfsl4KAGlmZWVzmLGvcCvcXtjZLWAI25eIiZmWUb1Bx4cWztwcCyIjSrOLPtVkk71XhOl6TlPqvNzKy+shu4pO2Bu4ALI2IDcAOwLzCRygg9edaK99SbmTVGVgOXtA2V5n1HRNwNEBGvR8THEbEZ+DmfcLiVmZnVX85RKKKyA2hNRFzTJ75bn7QT8OFWZmZNlXMUyhTgdOAZSd1F7BJgpqSJVE49fhk4pwH1mVnJTZgwITv3O9/5TlbeiSeemJU3duzY7Peut48//jgrr7e3d8Ccv/3tb8l4zlEojwNKPLQkETMzsybxmZhmZiXlBm5mVlJu4GZmJeUGbmZWUm7gZmYl5QZuZlZSbuBmZiXlBm5mVlKqsbJUY95MegP4S3F3F+CvTXvz5vHnap29IuKfW/HG/bbtLcrwb5ajUz4HlPezJLftpjbwf3hjaXknXqHQn8u26JR/s075HNBZnwU8hWJmVlpu4GZmJdXKBj6vhe/dSP5ctkWn/Jt1yueAzvosrZsDNzOz4fEUiplZSbmBm5mVVNMbuKTpkp6T9IKki5v9/vUk6VZJ6ySt7BPbWdJDkp4vvu/UyhqHQtKekh6VtEbSKkkXFPHSf7Zm6LBt/GVJz0jqlrS81fUMRqf+fvbV1AYuaRRwPXA0sD+VZdn2b2YNdTYfmN4vdjHwcESMBx4u7pfNJuCiiPg8cChwfvFz6oTP1lAduI0DHBERE0t4/PR8OvP38++aPQKfDLwQES9FxEfAr4Djm1xD3UTEUuCtfuHjgduK27cB32hmTfUQEb0RsaK4vRFYA+xOB3y2JuiobbzMOvX3s69mN/DdgVf73O8pYp1kTET0QqURAru2uJ5hkbQ3cDCwjA77bA3Sadt4AA9KelJSV6uLqYOO2oZzVqWvp9TiyD6OsU1J2h64C7gwIjZIqR+f9dNp2/iUiFgraVfgIUnPFiNbawPNHoH3AHv2ub8HsLbJNTTa65J2Ayi+r2txPUMiaRsqzfuOiLi7CHfEZ2uwjtrGI2Jt8X0dsIjKFFGZddQ23OwG/gQwXtJnJY0GTgMWN7mGRlsMnFncPhO4t4W1DIkqQ+1bgDURcU2fh0r/2ZqgY7ZxSdtJ2mHLbeAoYOUnP6vtddQ23PQzMSUdA1wLjAJujYjLm1pAHUlaCEylconK14EfAPcAvwbGAa8Ap0RE/x0pbU3SYcDvgGeAzUX4Eirz4KX+bM3QKdu4pH2ojLqhMt36yzJ9lk79/ezLp9KbmZWUz8Q0MyspN3Azs5JyAzczKyk3cDOzknIDNzMrKTdwM7OScgM3Myup/wOJLsMoRH+zFgAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "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": {},
   "source": [
    "We can then define a function that creates a basic digit classifier for the MNIST dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def create_model():\n",
    "    model = tf.keras.models.Sequential([\n",
    "        tf.keras.layers.Flatten(input_shape=(28, 28)),\n",
    "        tf.keras.layers.Dense(128, activation='relu'),\n",
    "        tf.keras.layers.Dropout(0.2),\n",
    "        tf.keras.layers.Dense(10)\n",
    "    ])\n",
    "\n",
    "    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss=loss_fn,\n",
    "                  metrics=['accuracy'])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now we move on with creating a model, opening the arrays, define a Tensorflow TileDB dataset and train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-08-10 16:27:34.953803: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  SSE4.1 SSE4.2 AVX AVX2 AVX512F FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n",
      "938/938 [==============================] - 5s 4ms/step - loss: 0.3426 - accuracy: 0.9022\n",
      "Epoch 2/5\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1658 - accuracy: 0.9510\n",
      "Epoch 3/5\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1239 - accuracy: 0.9632\n",
      "Epoch 4/5\n",
      "938/938 [==============================] - 4s 4ms/step - loss: 0.1015 - accuracy: 0.9688\n",
      "Epoch 5/5\n",
      "938/938 [==============================] - 5s 5ms/step - loss: 0.0866 - accuracy: 0.9732\n"
     ]
    }
   ],
   "source": [
    "from tiledb.ml.readers.tensorflow import TensorflowTileDBDataset, ArrayParams\n",
    "\n",
    "model = create_model()\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",
    "    tiledb_dataset = TensorflowTileDBDataset(\n",
    "        ArrayParams(array=x, fields=['features']),\n",
    "        ArrayParams(array=y, fields=['features']),\n",
    "        num_workers=2 if os.cpu_count() > 2 else 0,\n",
    "    )\n",
    "    tiledb_dataset = tiledb_dataset.batch(64).shuffle(128)\n",
    "    model.fit(tiledb_dataset, epochs=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " flatten (Flatten)           (None, 784)               0         \n",
      "                                                                 \n",
      " dense (Dense)               (None, 128)               100480    \n",
      "                                                                 \n",
      " dropout (Dropout)           (None, 128)               0         \n",
      "                                                                 \n",
      " dense_1 (Dense)             (None, 10)                1290      \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 101,770\n",
      "Trainable params: 101,770\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  }
 ],
 "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}