examples/models/pytorch_tiledb_models_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "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, and store it as TileDB array. Firstly, let's import what we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import glob\n",
    "import json\n",
    "import pickle\n",
    "import os\n",
    "import shutil\n",
    "from pprint import pprint\n",
    "\n",
    "import tiledb\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "\n",
    "from tiledb.ml.models.pytorch import PyTorchTileDBModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First let's define the parameters/hyperparameters we will need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<torch._C.Generator at 0x1193eb430>"
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "epochs = 1\n",
    "batch_size_train = 128\n",
    "learning_rate = 0.01\n",
    "momentum = 0.5\n",
    "log_interval = 10\n",
    "\n",
    "# Set random seeds for anything using random number generation\n",
    "random_seed = 1\n",
    "\n",
    "# Disable nondeterministic algorithms\n",
    "torch.backends.cudnn.enabled = False\n",
    "torch.manual_seed(random_seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We  will also need the DataLoaders API for the dataset. We will also employ TorchVision which let's as load the MNIST\n",
    "dataset in a handy way. We'll use a batch_size of 64 for training while the values 0.1307 and 0.3081 used for\n",
    "the Normalize() transformation below are the global mean and standard deviation of the MNIST dataset,\n",
    "we'll take them as a given here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "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",
      "\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"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100.0%\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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": [
      "15.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",
    "dataset = torchvision.datasets.MNIST(\n",
    "    root=data_home, \n",
    "    train=True, \n",
    "    download=True,\n",
    "    transform=torchvision.transforms.Compose([\n",
    "        torchvision.transforms.ToTensor(),\n",
    "        torchvision.transforms.Normalize((0.1307,), (0.3081,))\n",
    "    ])\n",
    ")\n",
    "train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size_train, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Moving on, we build our network. We'll use two 2-D convolutional layers followed by two fully-connected\n",
    "layers. As activation function we'll choose ReLUs and as a means of regularization we'll use two dropout layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
    "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
    "        self.conv2_drop = nn.Dropout2d()\n",
    "        self.fc1 = nn.Linear(320, 50)\n",
    "        self.fc2 = nn.Linear(50, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
    "        x = x.view(-1, 320)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim = 1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now initialise our Neural Network and optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "network = Net()\n",
    "optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We continue with the training loop, and we iterate over all training data once per epoch. Loading the individual batches\n",
    "is handled by the DataLoader. We need to set the gradients to zero using optimizer.zero_grad() since PyTorch by default\n",
    "accumulates gradients. We then produce the output of the network (forward pass) and compute a negative log-likelihodd\n",
    "loss between the output and the ground truth label. The backward() call we now collect a new set of gradients which we\n",
    "propagate back into each of the network's parameters using optimizer.step()."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 1 [0/60000 (0%)]\tLoss: 2.358812\n",
      "Train Epoch: 1 [1280/60000 (2%)]\tLoss: 2.285138\n",
      "Train Epoch: 1 [2560/60000 (4%)]\tLoss: 2.306635\n",
      "Train Epoch: 1 [3840/60000 (6%)]\tLoss: 2.270880\n",
      "Train Epoch: 1 [5120/60000 (9%)]\tLoss: 2.236740\n",
      "Train Epoch: 1 [6400/60000 (11%)]\tLoss: 2.243347\n",
      "Train Epoch: 1 [7680/60000 (13%)]\tLoss: 2.183255\n",
      "Train Epoch: 1 [8960/60000 (15%)]\tLoss: 2.148512\n",
      "Train Epoch: 1 [10240/60000 (17%)]\tLoss: 2.104912\n",
      "Train Epoch: 1 [11520/60000 (19%)]\tLoss: 2.004407\n",
      "Train Epoch: 1 [12800/60000 (21%)]\tLoss: 1.862252\n",
      "Train Epoch: 1 [14080/60000 (23%)]\tLoss: 1.884371\n",
      "Train Epoch: 1 [15360/60000 (26%)]\tLoss: 1.797316\n",
      "Train Epoch: 1 [16640/60000 (28%)]\tLoss: 1.687911\n",
      "Train Epoch: 1 [17920/60000 (30%)]\tLoss: 1.508046\n",
      "Train Epoch: 1 [19200/60000 (32%)]\tLoss: 1.764280\n",
      "Train Epoch: 1 [20480/60000 (34%)]\tLoss: 1.470073\n",
      "Train Epoch: 1 [21760/60000 (36%)]\tLoss: 1.351447\n",
      "Train Epoch: 1 [23040/60000 (38%)]\tLoss: 1.290582\n",
      "Train Epoch: 1 [24320/60000 (41%)]\tLoss: 1.017757\n",
      "Train Epoch: 1 [25600/60000 (43%)]\tLoss: 1.042162\n",
      "Train Epoch: 1 [26880/60000 (45%)]\tLoss: 1.098766\n",
      "Train Epoch: 1 [28160/60000 (47%)]\tLoss: 1.228552\n",
      "Train Epoch: 1 [29440/60000 (49%)]\tLoss: 1.149593\n",
      "Train Epoch: 1 [30720/60000 (51%)]\tLoss: 0.845248\n",
      "Train Epoch: 1 [32000/60000 (53%)]\tLoss: 0.974113\n",
      "Train Epoch: 1 [33280/60000 (55%)]\tLoss: 0.856906\n",
      "Train Epoch: 1 [34560/60000 (58%)]\tLoss: 0.923459\n",
      "Train Epoch: 1 [35840/60000 (60%)]\tLoss: 1.021857\n",
      "Train Epoch: 1 [37120/60000 (62%)]\tLoss: 0.806954\n",
      "Train Epoch: 1 [38400/60000 (64%)]\tLoss: 0.878951\n",
      "Train Epoch: 1 [39680/60000 (66%)]\tLoss: 0.818505\n",
      "Train Epoch: 1 [40960/60000 (68%)]\tLoss: 0.805543\n",
      "Train Epoch: 1 [42240/60000 (70%)]\tLoss: 0.823152\n",
      "Train Epoch: 1 [43520/60000 (72%)]\tLoss: 0.854361\n",
      "Train Epoch: 1 [44800/60000 (75%)]\tLoss: 0.774645\n",
      "Train Epoch: 1 [46080/60000 (77%)]\tLoss: 0.718349\n",
      "Train Epoch: 1 [47360/60000 (79%)]\tLoss: 0.543338\n",
      "Train Epoch: 1 [48640/60000 (81%)]\tLoss: 0.759377\n",
      "Train Epoch: 1 [49920/60000 (83%)]\tLoss: 0.654922\n",
      "Train Epoch: 1 [51200/60000 (85%)]\tLoss: 0.699930\n",
      "Train Epoch: 1 [52480/60000 (87%)]\tLoss: 0.805351\n",
      "Train Epoch: 1 [53760/60000 (90%)]\tLoss: 0.790734\n",
      "Train Epoch: 1 [55040/60000 (92%)]\tLoss: 0.759933\n",
      "Train Epoch: 1 [56320/60000 (94%)]\tLoss: 0.540410\n",
      "Train Epoch: 1 [57600/60000 (96%)]\tLoss: 0.641233\n",
      "Train Epoch: 1 [58880/60000 (98%)]\tLoss: 0.659374\n"
     ]
    }
   ],
   "source": [
    "train_losses = []\n",
    "train_counter = []\n",
    "\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "writer = SummaryWriter()\n",
    "\n",
    "\n",
    "def train(epoch):\n",
    "  network.train()\n",
    "  for batch_idx, (data, target) in enumerate(train_loader):\n",
    "    optimizer.zero_grad()\n",
    "    output = network(data)\n",
    "    loss = F.nll_loss(output, target)\n",
    "    writer.add_scalar(\"Loss/train\", loss, epoch)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    if batch_idx % log_interval == 0:\n",
    "      print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "        epoch, batch_idx * len(data), len(train_loader.dataset),\n",
    "        100. * batch_idx / len(train_loader), loss.item()))\n",
    "      train_losses.append(loss.item())\n",
    "      train_counter.append(\n",
    "        (batch_idx*64) + ((epoch-1)*len(train_loader.dataset)))\n",
    "\n",
    "for epoch in range(1, epochs + 1):\n",
    "  train(epoch)\n",
    "writer.flush()\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now save the trained model as a TileDB array. In case we want to train the model further in a later time, we can also save\n",
    "the optimizer in our TileDB array. In case we will use our model only for inference, we don't have to save the optimizer, and we\n",
    "only keep the model. We first declare a PytTorchTileDB object and initialize it with the corresponding TileDB uri, model and optimizer,\n",
    "and then save the model as a TileDB array. Finally, we can save any kind of metadata (in any structure, i.e., list, tuple or dictionary)\n",
    "by passing a dictionary to the meta attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/george/PycharmProjects/TileDB-ML/.venv/lib/python3.9/site-packages/tiledb/ctx.py:410: 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": [
    "uri = os.path.join(data_home, 'pytorch-mnist-1')\n",
    "tiledb_model_1 = PyTorchTileDBModel(uri=uri, model=network, optimizer=optimizer)\n",
    "\n",
    "tiledb_model_1.save(meta={'epochs': epochs,\n",
    "                          'train_loss': train_losses},\n",
    "                    summary_writer=writer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above step will create a TileDB array in your working directory. For information about the structure of a dense\n",
    "TileDB array in terms of files on disk please take a look [here](https://docs.tiledb.com/main/concepts/data-format).\n",
    "Let's open our TileDB array model and check metadata. Metadata that are of type list, dict or tuple have been JSON\n",
    "serialized while saving, i.e., we need json.loads to deserialize them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['../data/pytorch-mnist-1/__meta',\n",
      " '../data/pytorch-mnist-1/__fragment_meta',\n",
      " '../data/pytorch-mnist-1/__commits',\n",
      " '../data/pytorch-mnist-1/__schema',\n",
      " '../data/pytorch-mnist-1/__fragments']\n",
      "Key: TILEDB_ML_MODEL_ML_FRAMEWORK, Value: PYTORCH\n",
      "Key: TILEDB_ML_MODEL_ML_FRAMEWORK_VERSION, Value: 1.12.0\n",
      "Key: TILEDB_ML_MODEL_PREVIEW, Value: Net(\n",
      "  (conv1): Conv2d(1, 10, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (conv2): Conv2d(10, 20, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (conv2_drop): Dropout2d(p=0.5, inplace=False)\n",
      "  (fc1): Linear(in_features=320, out_features=50, bias=True)\n",
      "  (fc2): Linear(in_features=50, out_features=10, bias=True)\n",
      ")\n",
      "Key: TILEDB_ML_MODEL_PYTHON_VERSION, Value: 3.9.9\n",
      "Key: TILEDB_ML_MODEL_STAGE, Value: STAGING\n",
      "Key: TILEDB_ML_MODEL_VERSION, Value: \n",
      "Key: epochs, Value: 1\n",
      "Key: model_size, Value: 90053\n",
      "Key: optimizer_size, Value: 90064\n",
      "Key: tensorboard_size, Value: 22674\n",
      "Key: train_loss, Value: (2.358812093734741, 2.285137891769409, 2.3066349029541016, 2.2708795070648193, 2.2367401123046875, 2.24334716796875, 2.1832549571990967, 2.1485116481781006, 2.1049115657806396, 2.0044069290161133, 1.8622523546218872, 1.8843708038330078, 1.7973158359527588, 1.6879109144210815, 1.508046269416809, 1.764279842376709, 1.4700727462768555, 1.3514467477798462, 1.2905819416046143, 1.0177571773529053, 1.042162299156189, 1.0987662076950073, 1.2285516262054443, 1.1495932340621948, 0.8452475070953369, 0.9741130471229553, 0.8569056987762451, 0.9234588146209717, 1.0218565464019775, 0.8069543242454529, 0.8789511919021606, 0.8185049891471863, 0.8055434226989746, 0.8231522440910339, 0.8543609976768494, 0.7746452689170837, 0.718348503112793, 0.5433375239372253, 0.7593768239021301, 0.65492182970047, 0.6999298930168152, 0.8053513765335083, 0.790733814239502, 0.7599329948425293, 0.540409505367279, 0.6412327885627747, 0.6593738198280334)\n"
     ]
    }
   ],
   "source": [
    "# Check array directory\n",
    "pprint(glob.glob(f'{uri}/*'))\n",
    "\n",
    "# Open in write mode in order to add metadata\n",
    "model_array_1 = tiledb.open(uri)\n",
    "for key, value in model_array_1.meta.items():\n",
    "    if isinstance(value, bytes):\n",
    "        value = json.loads(value)\n",
    "    print(\"Key: {}, Value: {}\".format(key, value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, in array's metadata we have by default information about the backend we used for training (pytorch),\n",
    "pytorch version, python version and the extra metadata about epochs and training loss that we added.\n",
    "We can load and check any of the aforementioned without having to load the entire model in memory.\n",
    "Moreover, we can add any kind of extra information in model's metadata also by opening the TileDB array and adding new keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Key: TILEDB_ML_MODEL_ML_FRAMEWORK, Value: PYTORCH\n",
      "Key: TILEDB_ML_MODEL_ML_FRAMEWORK_VERSION, Value: 1.12.0\n",
      "Key: TILEDB_ML_MODEL_PREVIEW, Value: Net(\n",
      "  (conv1): Conv2d(1, 10, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (conv2): Conv2d(10, 20, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (conv2_drop): Dropout2d(p=0.5, inplace=False)\n",
      "  (fc1): Linear(in_features=320, out_features=50, bias=True)\n",
      "  (fc2): Linear(in_features=50, out_features=10, bias=True)\n",
      ")\n",
      "Key: TILEDB_ML_MODEL_PYTHON_VERSION, Value: 3.9.9\n",
      "Key: TILEDB_ML_MODEL_STAGE, Value: STAGING\n",
      "Key: TILEDB_ML_MODEL_VERSION, Value: \n",
      "Key: epochs, Value: 1\n",
      "Key: model_size, Value: 90053\n",
      "Key: new_meta, Value: [\"Any kind of info\"]\n",
      "Key: optimizer_size, Value: 90064\n",
      "Key: tensorboard_size, Value: 22674\n",
      "Key: train_loss, Value: (2.358812093734741, 2.285137891769409, 2.3066349029541016, 2.2708795070648193, 2.2367401123046875, 2.24334716796875, 2.1832549571990967, 2.1485116481781006, 2.1049115657806396, 2.0044069290161133, 1.8622523546218872, 1.8843708038330078, 1.7973158359527588, 1.6879109144210815, 1.508046269416809, 1.764279842376709, 1.4700727462768555, 1.3514467477798462, 1.2905819416046143, 1.0177571773529053, 1.042162299156189, 1.0987662076950073, 1.2285516262054443, 1.1495932340621948, 0.8452475070953369, 0.9741130471229553, 0.8569056987762451, 0.9234588146209717, 1.0218565464019775, 0.8069543242454529, 0.8789511919021606, 0.8185049891471863, 0.8055434226989746, 0.8231522440910339, 0.8543609976768494, 0.7746452689170837, 0.718348503112793, 0.5433375239372253, 0.7593768239021301, 0.65492182970047, 0.6999298930168152, 0.8053513765335083, 0.790733814239502, 0.7599329948425293, 0.540409505367279, 0.6412327885627747, 0.6593738198280334)\n"
     ]
    }
   ],
   "source": [
    "# Open the array in write mode\n",
    "with tiledb.Array(uri, \"w\") as A:\n",
    "    # Keep all history\n",
    "    A.meta['new_meta'] = json.dumps(['Any kind of info'])\n",
    "\n",
    "# Check that everything is there\n",
    "model_array_1 = tiledb.open(uri)\n",
    "for key, value in model_array_1.meta.items():\n",
    "    if isinstance(value, bytes):\n",
    "        value = json.loads(value)\n",
    "    print(\"Key: {}, Value: {}\".format(key, value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Moving on, we can load the trained models for prediction, evaluation or retraining, as usual with\n",
    "PyTorch models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Place holder for the loaded model\n",
    "network = Net()\n",
    "optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)\n",
    "tiledb_model_1.load(model=network, optimizer=optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What is really nice with saving models as TileDB array, is native versioning based on fragments as described [here](https://docs.tiledb.com/main/concepts/data-format#immutable-fragments). We can load a model, retrain it with new data and update the already existing TileDB model array with new model parameters and metadata. All information, old and new will be there and accessible. This is extremely useful when you retrain with new data or trying different architectures for the same problem, and you want to keep track of all your experiments without having to store different model instances. In our case, let's continue training model_1 with the rest of our dataset and for 2 more epochs. After training is done, you will notice the extra directories and files (fragments) added to `pytorch-mnist-1` TileDB array directory, which keep all versions of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 1 [0/60000 (0%)]\tLoss: 0.576272\n",
      "Train Epoch: 1 [1280/60000 (2%)]\tLoss: 0.600712\n",
      "Train Epoch: 1 [2560/60000 (4%)]\tLoss: 0.750536\n",
      "Train Epoch: 1 [3840/60000 (6%)]\tLoss: 0.668426\n",
      "Train Epoch: 1 [5120/60000 (9%)]\tLoss: 0.645401\n",
      "Train Epoch: 1 [6400/60000 (11%)]\tLoss: 0.544283\n",
      "Train Epoch: 1 [7680/60000 (13%)]\tLoss: 0.664902\n",
      "Train Epoch: 1 [8960/60000 (15%)]\tLoss: 0.463885\n",
      "Train Epoch: 1 [10240/60000 (17%)]\tLoss: 0.633202\n",
      "Train Epoch: 1 [11520/60000 (19%)]\tLoss: 0.620702\n",
      "Train Epoch: 1 [12800/60000 (21%)]\tLoss: 0.638514\n",
      "Train Epoch: 1 [14080/60000 (23%)]\tLoss: 0.521937\n",
      "Train Epoch: 1 [15360/60000 (26%)]\tLoss: 0.568641\n",
      "Train Epoch: 1 [16640/60000 (28%)]\tLoss: 0.403443\n",
      "Train Epoch: 1 [17920/60000 (30%)]\tLoss: 0.734374\n",
      "Train Epoch: 1 [19200/60000 (32%)]\tLoss: 0.519801\n",
      "Train Epoch: 1 [20480/60000 (34%)]\tLoss: 0.616732\n",
      "Train Epoch: 1 [21760/60000 (36%)]\tLoss: 0.627395\n",
      "Train Epoch: 1 [23040/60000 (38%)]\tLoss: 0.622394\n",
      "Train Epoch: 1 [24320/60000 (41%)]\tLoss: 0.427548\n",
      "Train Epoch: 1 [25600/60000 (43%)]\tLoss: 0.572651\n",
      "Train Epoch: 1 [26880/60000 (45%)]\tLoss: 0.436058\n",
      "Train Epoch: 1 [28160/60000 (47%)]\tLoss: 0.502791\n",
      "Train Epoch: 1 [29440/60000 (49%)]\tLoss: 0.477583\n",
      "Train Epoch: 1 [30720/60000 (51%)]\tLoss: 0.555044\n",
      "Train Epoch: 1 [32000/60000 (53%)]\tLoss: 0.510869\n",
      "Train Epoch: 1 [33280/60000 (55%)]\tLoss: 0.548664\n",
      "Train Epoch: 1 [34560/60000 (58%)]\tLoss: 0.530004\n",
      "Train Epoch: 1 [35840/60000 (60%)]\tLoss: 0.411406\n",
      "Train Epoch: 1 [37120/60000 (62%)]\tLoss: 0.554757\n",
      "Train Epoch: 1 [38400/60000 (64%)]\tLoss: 0.505362\n",
      "Train Epoch: 1 [39680/60000 (66%)]\tLoss: 0.437215\n",
      "Train Epoch: 1 [40960/60000 (68%)]\tLoss: 0.423777\n",
      "Train Epoch: 1 [42240/60000 (70%)]\tLoss: 0.358317\n",
      "Train Epoch: 1 [43520/60000 (72%)]\tLoss: 0.378239\n",
      "Train Epoch: 1 [44800/60000 (75%)]\tLoss: 0.427609\n",
      "Train Epoch: 1 [46080/60000 (77%)]\tLoss: 0.643378\n",
      "Train Epoch: 1 [47360/60000 (79%)]\tLoss: 0.595157\n",
      "Train Epoch: 1 [48640/60000 (81%)]\tLoss: 0.643645\n",
      "Train Epoch: 1 [49920/60000 (83%)]\tLoss: 0.420812\n",
      "Train Epoch: 1 [51200/60000 (85%)]\tLoss: 0.633425\n",
      "Train Epoch: 1 [52480/60000 (87%)]\tLoss: 0.566453\n",
      "Train Epoch: 1 [53760/60000 (90%)]\tLoss: 0.431015\n",
      "Train Epoch: 1 [55040/60000 (92%)]\tLoss: 0.635182\n",
      "Train Epoch: 1 [56320/60000 (94%)]\tLoss: 0.452681\n",
      "Train Epoch: 1 [57600/60000 (96%)]\tLoss: 0.506694\n",
      "Train Epoch: 1 [58880/60000 (98%)]\tLoss: 0.671234\n",
      "Train Epoch: 2 [0/60000 (0%)]\tLoss: 0.426433\n",
      "Train Epoch: 2 [1280/60000 (2%)]\tLoss: 0.542953\n",
      "Train Epoch: 2 [2560/60000 (4%)]\tLoss: 0.431448\n",
      "Train Epoch: 2 [3840/60000 (6%)]\tLoss: 0.333652\n",
      "Train Epoch: 2 [5120/60000 (9%)]\tLoss: 0.367727\n",
      "Train Epoch: 2 [6400/60000 (11%)]\tLoss: 0.578225\n",
      "Train Epoch: 2 [7680/60000 (13%)]\tLoss: 0.576888\n",
      "Train Epoch: 2 [8960/60000 (15%)]\tLoss: 0.269501\n",
      "Train Epoch: 2 [10240/60000 (17%)]\tLoss: 0.382370\n",
      "Train Epoch: 2 [11520/60000 (19%)]\tLoss: 0.395886\n",
      "Train Epoch: 2 [12800/60000 (21%)]\tLoss: 0.374322\n",
      "Train Epoch: 2 [14080/60000 (23%)]\tLoss: 0.297448\n",
      "Train Epoch: 2 [15360/60000 (26%)]\tLoss: 0.475420\n",
      "Train Epoch: 2 [16640/60000 (28%)]\tLoss: 0.296931\n",
      "Train Epoch: 2 [17920/60000 (30%)]\tLoss: 0.549110\n",
      "Train Epoch: 2 [19200/60000 (32%)]\tLoss: 0.354616\n",
      "Train Epoch: 2 [20480/60000 (34%)]\tLoss: 0.432848\n",
      "Train Epoch: 2 [21760/60000 (36%)]\tLoss: 0.382388\n",
      "Train Epoch: 2 [23040/60000 (38%)]\tLoss: 0.276277\n",
      "Train Epoch: 2 [24320/60000 (41%)]\tLoss: 0.454842\n",
      "Train Epoch: 2 [25600/60000 (43%)]\tLoss: 0.353321\n",
      "Train Epoch: 2 [26880/60000 (45%)]\tLoss: 0.451096\n",
      "Train Epoch: 2 [28160/60000 (47%)]\tLoss: 0.316580\n",
      "Train Epoch: 2 [29440/60000 (49%)]\tLoss: 0.364130\n",
      "Train Epoch: 2 [30720/60000 (51%)]\tLoss: 0.431465\n",
      "Train Epoch: 2 [32000/60000 (53%)]\tLoss: 0.524982\n",
      "Train Epoch: 2 [33280/60000 (55%)]\tLoss: 0.438477\n",
      "Train Epoch: 2 [34560/60000 (58%)]\tLoss: 0.454826\n",
      "Train Epoch: 2 [35840/60000 (60%)]\tLoss: 0.518262\n",
      "Train Epoch: 2 [37120/60000 (62%)]\tLoss: 0.471013\n",
      "Train Epoch: 2 [38400/60000 (64%)]\tLoss: 0.274934\n",
      "Train Epoch: 2 [39680/60000 (66%)]\tLoss: 0.398929\n",
      "Train Epoch: 2 [40960/60000 (68%)]\tLoss: 0.395629\n",
      "Train Epoch: 2 [42240/60000 (70%)]\tLoss: 0.404211\n",
      "Train Epoch: 2 [43520/60000 (72%)]\tLoss: 0.466308\n",
      "Train Epoch: 2 [44800/60000 (75%)]\tLoss: 0.306582\n",
      "Train Epoch: 2 [46080/60000 (77%)]\tLoss: 0.306776\n",
      "Train Epoch: 2 [47360/60000 (79%)]\tLoss: 0.521016\n",
      "Train Epoch: 2 [48640/60000 (81%)]\tLoss: 0.333933\n",
      "Train Epoch: 2 [49920/60000 (83%)]\tLoss: 0.384981\n",
      "Train Epoch: 2 [51200/60000 (85%)]\tLoss: 0.307746\n",
      "Train Epoch: 2 [52480/60000 (87%)]\tLoss: 0.305437\n",
      "Train Epoch: 2 [53760/60000 (90%)]\tLoss: 0.244203\n",
      "Train Epoch: 2 [55040/60000 (92%)]\tLoss: 0.301768\n",
      "Train Epoch: 2 [56320/60000 (94%)]\tLoss: 0.211211\n",
      "Train Epoch: 2 [57600/60000 (96%)]\tLoss: 0.319129\n",
      "Train Epoch: 2 [58880/60000 (98%)]\tLoss: 0.281229\n",
      "\n",
      "['../data/pytorch-mnist-1/__meta',\n",
      " '../data/pytorch-mnist-1/__fragment_meta',\n",
      " '../data/pytorch-mnist-1/__commits',\n",
      " '../data/pytorch-mnist-1/__schema',\n",
      " '../data/pytorch-mnist-1/__fragments']\n",
      "\n",
      "====== FRAGMENTS  INFO ======\n",
      "array uri: ../data/pytorch-mnist-1\n",
      "number of fragments: 2\n",
      "\n",
      "===== FRAGMENT NUMBER 0 =====\n",
      "fragment uri: file:///Users/george/PycharmProjects/TileDB-ML/examples/data/pytorch-mnist-1/__fragments/__1675169603990_1675169603990_aae704a1499649fe8ad3fd0e61d8f9b9_16\n",
      "timestamp range: (1675169603990, 1675169603990)\n",
      "number of unconsolidated metadata: 2\n",
      "version: 16\n",
      "\n",
      "===== FRAGMENT NUMBER 1 =====\n",
      "fragment uri: file:///Users/george/PycharmProjects/TileDB-ML/examples/data/pytorch-mnist-1/__fragments/__1675169635431_1675169635431_e064ef7a982e45a1be7bb678d1949b97_16\n",
      "timestamp range: (1675169635431, 1675169635431)\n",
      "number of unconsolidated metadata: 2\n",
      "version: 16\n"
     ]
    }
   ],
   "source": [
    "train_losses = []\n",
    "train_counter = []\n",
    "\n",
    "# We train for some extra 2 epochs\n",
    "for epoch in range(1, 2 + 1):\n",
    "  train(epoch)\n",
    "\n",
    "# and update\n",
    "tiledb_model_1 = PyTorchTileDBModel(uri=uri, model=network, optimizer=optimizer)\n",
    "tiledb_model_1.save(meta={'epochs': epochs,\n",
    "                          'train_loss': train_losses})\n",
    "\n",
    "# Check array directory\n",
    "print()\n",
    "pprint(glob.glob(f'{uri}/*'))\n",
    "\n",
    "# tiledb.array_fragments() requires TileDB-Py version > 0.8.5\n",
    "fragments_info = tiledb.array_fragments(uri)\n",
    "\n",
    "print()\n",
    "print(\"====== FRAGMENTS  INFO ======\")\n",
    "print(\"array uri: {}\".format(fragments_info.array_uri))\n",
    "print(\"number of fragments: {}\".format(len(fragments_info)))\n",
    "\n",
    "for fragment_num, fragment in enumerate(fragments_info, start=1):\n",
    "    print()\n",
    "    print(\"===== FRAGMENT NUMBER {} =====\".format(fragment.num))\n",
    "    print(\"fragment uri: {}\".format(fragment.uri))\n",
    "    print(\"timestamp range: {}\".format(fragment.timestamp_range))\n",
    "    print(\n",
    "        \"number of unconsolidated metadata: {}\".format(\n",
    "            fragment.unconsolidated_metadata_num\n",
    "        )\n",
    "    )\n",
    "    print(\"version: {}\".format(fragment.version))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, a very interesting and useful, for machine learning models, TileDB feature that is described\n",
    "[here](https://docs.tiledb.com/main/concepts/data-format#groups) and [here](https://docs.tiledb.com/main/how-to/object-management#creating-tiledb-groups)\n",
    "are groups. Assuming we want to solve the MNIST problem, and we want to try several architectures. We can save each architecture\n",
    "as a separate TileDB array with native versioning each time it is re-trained, and then organise all models that solve the same problem (MNIST)\n",
    "as a TileDB array group with any kind of hierarchy. Let's firstly define a new model architecture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class OtherNet(nn.Module):\n",
    "    # For the sake of simplicity we just tweak the initial architecture by replacing a relu with relu6.\n",
    "    def __init__(self):\n",
    "        super(OtherNet, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
    "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
    "        self.conv2_drop = nn.Dropout2d()\n",
    "        self.fc1 = nn.Linear(320, 50)\n",
    "        self.fc2 = nn.Linear(50, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
    "        x = x.view(-1, 320)\n",
    "        x = F.relu6(self.fc1(x))\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim = 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then train it and save it as a new TileDB array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 1 [0/60000 (0%)]\tLoss: 2.313390\n",
      "Train Epoch: 1 [1280/60000 (2%)]\tLoss: 1.527101\n",
      "Train Epoch: 1 [2560/60000 (4%)]\tLoss: 0.971071\n",
      "Train Epoch: 1 [3840/60000 (6%)]\tLoss: 0.857200\n",
      "Train Epoch: 1 [5120/60000 (9%)]\tLoss: 0.513735\n",
      "Train Epoch: 1 [6400/60000 (11%)]\tLoss: 0.778414\n",
      "Train Epoch: 1 [7680/60000 (13%)]\tLoss: 0.419221\n",
      "Train Epoch: 1 [8960/60000 (15%)]\tLoss: 0.507100\n",
      "Train Epoch: 1 [10240/60000 (17%)]\tLoss: 0.543142\n",
      "Train Epoch: 1 [11520/60000 (19%)]\tLoss: 0.465935\n",
      "Train Epoch: 1 [12800/60000 (21%)]\tLoss: 0.572005\n",
      "Train Epoch: 1 [14080/60000 (23%)]\tLoss: 0.583393\n",
      "Train Epoch: 1 [15360/60000 (26%)]\tLoss: 0.410894\n",
      "Train Epoch: 1 [16640/60000 (28%)]\tLoss: 0.496960\n",
      "Train Epoch: 1 [17920/60000 (30%)]\tLoss: 0.500474\n",
      "Train Epoch: 1 [19200/60000 (32%)]\tLoss: 0.343156\n",
      "Train Epoch: 1 [20480/60000 (34%)]\tLoss: 0.501805\n",
      "Train Epoch: 1 [21760/60000 (36%)]\tLoss: 0.543956\n",
      "Train Epoch: 1 [23040/60000 (38%)]\tLoss: 0.558721\n",
      "Train Epoch: 1 [24320/60000 (41%)]\tLoss: 0.365339\n",
      "Train Epoch: 1 [25600/60000 (43%)]\tLoss: 0.427882\n",
      "Train Epoch: 1 [26880/60000 (45%)]\tLoss: 0.413578\n",
      "Train Epoch: 1 [28160/60000 (47%)]\tLoss: 0.437229\n",
      "Train Epoch: 1 [29440/60000 (49%)]\tLoss: 0.424080\n",
      "Train Epoch: 1 [30720/60000 (51%)]\tLoss: 0.486696\n",
      "Train Epoch: 1 [32000/60000 (53%)]\tLoss: 0.345004\n",
      "Train Epoch: 1 [33280/60000 (55%)]\tLoss: 0.429583\n",
      "Train Epoch: 1 [34560/60000 (58%)]\tLoss: 0.329117\n",
      "Train Epoch: 1 [35840/60000 (60%)]\tLoss: 0.303430\n",
      "Train Epoch: 1 [37120/60000 (62%)]\tLoss: 0.379547\n",
      "Train Epoch: 1 [38400/60000 (64%)]\tLoss: 0.465007\n",
      "Train Epoch: 1 [39680/60000 (66%)]\tLoss: 0.301131\n",
      "Train Epoch: 1 [40960/60000 (68%)]\tLoss: 0.264844\n",
      "Train Epoch: 1 [42240/60000 (70%)]\tLoss: 0.346164\n",
      "Train Epoch: 1 [43520/60000 (72%)]\tLoss: 0.441803\n",
      "Train Epoch: 1 [44800/60000 (75%)]\tLoss: 0.318938\n",
      "Train Epoch: 1 [46080/60000 (77%)]\tLoss: 0.254888\n",
      "Train Epoch: 1 [47360/60000 (79%)]\tLoss: 0.319974\n",
      "Train Epoch: 1 [48640/60000 (81%)]\tLoss: 0.382300\n",
      "Train Epoch: 1 [49920/60000 (83%)]\tLoss: 0.404844\n",
      "Train Epoch: 1 [51200/60000 (85%)]\tLoss: 0.267870\n",
      "Train Epoch: 1 [52480/60000 (87%)]\tLoss: 0.374096\n",
      "Train Epoch: 1 [53760/60000 (90%)]\tLoss: 0.442077\n",
      "Train Epoch: 1 [55040/60000 (92%)]\tLoss: 0.533317\n",
      "Train Epoch: 1 [56320/60000 (94%)]\tLoss: 0.429747\n",
      "Train Epoch: 1 [57600/60000 (96%)]\tLoss: 0.439626\n",
      "Train Epoch: 1 [58880/60000 (98%)]\tLoss: 0.338414\n"
     ]
    }
   ],
   "source": [
    "network = OtherNet()\n",
    "optimizer = optim.Adam(network.parameters(), lr=learning_rate)\n",
    "\n",
    "train_losses = []\n",
    "train_counter = []\n",
    "\n",
    "for epoch in range(1, epochs + 1):\n",
    "    train(epoch)\n",
    "\n",
    "uri2 = os.path.join(data_home, 'pytorch-mnist-2')\n",
    "tiledb_model_2 = PyTorchTileDBModel(uri=uri2, model=network, optimizer=optimizer)\n",
    "\n",
    "tiledb_model_2.save(meta={'epochs': epochs,\n",
    "                          'train_loss': train_losses})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create a TileDB group and organise (in hierarchies, e.g., sophisticated vs less sophisticated) all our\n",
    "MNIST models as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "'../data/tiledb-pytorch-mnist/pytorch-mnist-2'"
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "group = os.path.join(data_home, 'tiledb-pytorch-mnist')\n",
    "tiledb.group_create(group)\n",
    "shutil.move(uri, group)\n",
    "shutil.move(uri2, group)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Any time we can check and query all the available models, including their metadata, for a specific problem like MNIST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "file:///Users/george/PycharmProjects/TileDB-ML/examples/data/tiledb-pytorch-mnist/pytorch-mnist-1 array\n",
      "file:///Users/george/PycharmProjects/TileDB-ML/examples/data/tiledb-pytorch-mnist/pytorch-mnist-2 array\n"
     ]
    }
   ],
   "source": [
    "tiledb.ls(group, lambda obj_path, obj_type: print(obj_path, obj_type))"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}