examples/cloud/models/pytorch_tiledb_cloud_ml_model_array.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we show how we can train a model with PyTorch and save it as a TileDB array on TileDB-Cloud.\n",
    "Firstly, let's import what we need and define some variables needed for training a model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "is_executing": true,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import tiledb.cloud\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\n",
    "\n",
    "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",
    "torch.manual_seed(seed=1)\n",
    "\n",
    "# Disable nondeterministic algorithms\n",
    "torch.backends.cudnn.enabled = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then create a TileDB-Cloud context, log in to TileDB via our API token (or username/password) and get our username."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ctx = tiledb.cloud.Ctx()\n",
    "tiledb.cloud.login(token=os.getenv('TILEDB_API_TOKEN'))\n",
    "namespace = tiledb.cloud.client.default_user().username"
   ]
  },
  {
   "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": [],
   "source": [
    "data_home = os.path.join(os.path.pardir, 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": [
    "model = Net()\n",
    "optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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.236739\n",
      "Train Epoch: 1 [6400/60000 (11%)]\tLoss: 2.243352\n",
      "Train Epoch: 1 [7680/60000 (13%)]\tLoss: 2.183357\n",
      "Train Epoch: 1 [8960/60000 (15%)]\tLoss: 2.148380\n",
      "Train Epoch: 1 [10240/60000 (17%)]\tLoss: 2.104831\n",
      "Train Epoch: 1 [11520/60000 (19%)]\tLoss: 2.004480\n",
      "Train Epoch: 1 [12800/60000 (21%)]\tLoss: 1.861931\n",
      "Train Epoch: 1 [14080/60000 (23%)]\tLoss: 1.884068\n",
      "Train Epoch: 1 [15360/60000 (26%)]\tLoss: 1.797761\n",
      "Train Epoch: 1 [16640/60000 (28%)]\tLoss: 1.687489\n",
      "Train Epoch: 1 [17920/60000 (30%)]\tLoss: 1.509098\n",
      "Train Epoch: 1 [19200/60000 (32%)]\tLoss: 1.766841\n",
      "Train Epoch: 1 [20480/60000 (34%)]\tLoss: 1.469026\n",
      "Train Epoch: 1 [21760/60000 (36%)]\tLoss: 1.349417\n",
      "Train Epoch: 1 [23040/60000 (38%)]\tLoss: 1.291566\n",
      "Train Epoch: 1 [24320/60000 (41%)]\tLoss: 1.022800\n",
      "Train Epoch: 1 [25600/60000 (43%)]\tLoss: 1.045352\n",
      "Train Epoch: 1 [26880/60000 (45%)]\tLoss: 1.095343\n",
      "Train Epoch: 1 [28160/60000 (47%)]\tLoss: 1.226705\n",
      "Train Epoch: 1 [29440/60000 (49%)]\tLoss: 1.148985\n",
      "Train Epoch: 1 [30720/60000 (51%)]\tLoss: 0.845941\n",
      "Train Epoch: 1 [32000/60000 (53%)]\tLoss: 0.968638\n",
      "Train Epoch: 1 [33280/60000 (55%)]\tLoss: 0.853761\n",
      "Train Epoch: 1 [34560/60000 (58%)]\tLoss: 0.926097\n",
      "Train Epoch: 1 [35840/60000 (60%)]\tLoss: 1.018820\n",
      "Train Epoch: 1 [37120/60000 (62%)]\tLoss: 0.807615\n",
      "Train Epoch: 1 [38400/60000 (64%)]\tLoss: 0.879665\n",
      "Train Epoch: 1 [39680/60000 (66%)]\tLoss: 0.820776\n",
      "Train Epoch: 1 [40960/60000 (68%)]\tLoss: 0.801570\n",
      "Train Epoch: 1 [42240/60000 (70%)]\tLoss: 0.825205\n",
      "Train Epoch: 1 [43520/60000 (72%)]\tLoss: 0.855643\n",
      "Train Epoch: 1 [44800/60000 (75%)]\tLoss: 0.768875\n",
      "Train Epoch: 1 [46080/60000 (77%)]\tLoss: 0.722621\n",
      "Train Epoch: 1 [47360/60000 (79%)]\tLoss: 0.548631\n",
      "Train Epoch: 1 [48640/60000 (81%)]\tLoss: 0.760798\n",
      "Train Epoch: 1 [49920/60000 (83%)]\tLoss: 0.654459\n",
      "Train Epoch: 1 [51200/60000 (85%)]\tLoss: 0.697698\n",
      "Train Epoch: 1 [52480/60000 (87%)]\tLoss: 0.805110\n",
      "Train Epoch: 1 [53760/60000 (90%)]\tLoss: 0.790563\n",
      "Train Epoch: 1 [55040/60000 (92%)]\tLoss: 0.759908\n",
      "Train Epoch: 1 [56320/60000 (94%)]\tLoss: 0.535904\n",
      "Train Epoch: 1 [57600/60000 (96%)]\tLoss: 0.643054\n",
      "Train Epoch: 1 [58880/60000 (98%)]\tLoss: 0.657785\n"
     ]
    }
   ],
   "source": [
    "train_losses = []\n",
    "train_counter = []\n",
    "\n",
    "def train(epoch):\n",
    "  model.train()\n",
    "  for batch_idx, (data, target) in enumerate(train_loader):\n",
    "    optimizer.zero_grad()\n",
    "    output = model(data)\n",
    "    loss = F.nll_loss(output, target)\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can move on by defining a TileDBPyTorch model and use model save functionality in order to save it directly to\n",
    "our bucket on S3 (defined with AWS credentials in your TileDB-Cloud account) and register it on TileDB-Cloud."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defining PyTorchTileDBModel model...\n",
      "Saving model on S3 and registering on TileDB-Cloud...\n"
     ]
    }
   ],
   "source": [
    "print('Defining PyTorchTileDBModel model...')\n",
    "# In order to save our model on S3 and register it on TileDB-Cloud we have to pass our Namespace and TileDB Context.\n",
    "tiledb_model = PyTorchTileDBModel(uri='tiledb-pytorch-model', namespace=namespace, ctx=ctx, model=model, optimizer=optimizer)\n",
    "\n",
    "# We will need the uri that was created from our model class\n",
    "# (and follows pattern tiledb://my_username/s3://my_bucket/my_array),\n",
    "# in order to interact with our model on TileDB-Cloud.\n",
    "tiledb_cloud_model_uri = tiledb_model.uri\n",
    "\n",
    "print('Saving model on S3 and registering on TileDB-Cloud...')\n",
    "tiledb_model.save(meta={'epochs': epochs,\n",
    "                        'train_loss': train_losses})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can use TileDB-Cloud API as described in our [cloud documentation](https://docs.tiledb.com/cloud/), in order\n",
    "to list our models, get information and deregister them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "is_executing": true,
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'arrays': [{'access_credentials_name': 'gsk',\n",
      "             'allowed_actions': ['read',\n",
      "                                 'read_array_logs',\n",
      "                                 'read_array_info',\n",
      "                                 'read_array_schema',\n",
      "                                 'write',\n",
      "                                 'edit'],\n",
      "             'description': None,\n",
      "             'file_properties': None,\n",
      "             'file_type': 'ml_model',\n",
      "             'id': '686687e4-f21f-42b6-90fc-464acf221a06',\n",
      "             'is_favorite': None,\n",
      "             'last_accessed': datetime.datetime(2022, 8, 10, 18, 9, 57, tzinfo=tzutc()),\n",
      "             'license_id': None,\n",
      "             'license_text': None,\n",
      "             'logo': None,\n",
      "             'name': 'tiledb-pytorch-model',\n",
      "             'namespace': 'george.sakkis',\n",
      "             'namespace_subscribed': None,\n",
      "             'pricing': None,\n",
      "             'public_share': False,\n",
      "             'read_only': False,\n",
      "             'share_count': None,\n",
      "             'size': 846577.0,\n",
      "             'subscriptions': None,\n",
      "             'tags': None,\n",
      "             'tiledb_uri': 'tiledb://george.sakkis/tiledb-pytorch-model',\n",
      "             'type': 'dense',\n",
      "             'uri': 's3://tiledb-gsk/ml_models/tiledb-pytorch-model'}],\n",
      " 'pagination_metadata': {'page': None,\n",
      "                         'per_page': None,\n",
      "                         'total_items': None,\n",
      "                         'total_pages': None}}\n",
      "{'access_credentials_name': 'gsk',\n",
      " 'allowed_actions': ['edit',\n",
      "                     'read',\n",
      "                     'read_array_info',\n",
      "                     'read_array_logs',\n",
      "                     'read_array_schema',\n",
      "                     'write'],\n",
      " 'description': None,\n",
      " 'file_properties': {'TILEDB_ML_MODEL_ML_FRAMEWORK': 'PYTORCH',\n",
      "                     'TILEDB_ML_MODEL_ML_FRAMEWORK_VERSION': '1.12.0.post2',\n",
      "                     'TILEDB_ML_MODEL_PREVIEW': 'Net(\\n'\n",
      "                                                '  (conv1): Conv2d(1, 10, '\n",
      "                                                'kernel_size=(5, 5), '\n",
      "                                                'stride=(1, 1))\\n'\n",
      "                                                '  (conv2): Conv2d(10, 20, '\n",
      "                                                'kernel_size=(5, 5), '\n",
      "                                                'stride=(1, 1))\\n'\n",
      "                                                '  (conv2_drop): '\n",
      "                                                'Dropout2d(p=0.5, '\n",
      "                                                'inplace=False)\\n'\n",
      "                                                '  (fc1): '\n",
      "                                                'Linear(in_features=320, '\n",
      "                                                'out_features=50, bias=True)\\n'\n",
      "                                                '  (fc2): '\n",
      "                                                'Linear(in_features=50, '\n",
      "                                                'out_features=10, bias=True)\\n'\n",
      "                                                ')',\n",
      "                     'TILEDB_ML_MODEL_PYTHON_VERSION': '3.7.13',\n",
      "                     'TILEDB_ML_MODEL_STAGE': 'STAGING'},\n",
      " 'file_type': 'ml_model',\n",
      " 'id': '686687e4-f21f-42b6-90fc-464acf221a06',\n",
      " 'is_favorite': None,\n",
      " 'last_accessed': datetime.datetime(2022, 8, 10, 18, 9, 57, tzinfo=tzutc()),\n",
      " 'license_id': None,\n",
      " 'license_text': None,\n",
      " 'logo': None,\n",
      " 'name': 'tiledb-pytorch-model',\n",
      " 'namespace': 'george.sakkis',\n",
      " 'namespace_subscribed': False,\n",
      " 'pricing': None,\n",
      " 'public_share': False,\n",
      " 'read_only': False,\n",
      " 'share_count': None,\n",
      " 'size': 846577.0,\n",
      " 'subscriptions': None,\n",
      " 'tags': None,\n",
      " 'tiledb_uri': 'tiledb://george.sakkis/tiledb-pytorch-model',\n",
      " 'type': 'dense',\n",
      " 'uri': 's3://tiledb-gsk/ml_models/tiledb-pytorch-model'}\n"
     ]
    }
   ],
   "source": [
    "# List all our models. Here, we filter with file_type = 'ml_model'. All machine learning model TileDB arrays are of type\n",
    "# 'ml_model'\n",
    "print(tiledb.cloud.client.list_arrays(file_type=['ml_model'], namespace=namespace))\n",
    "\n",
    "# Get model's info\n",
    "print(tiledb.cloud.array.info(tiledb_cloud_model_uri))\n",
    "\n",
    "# Load our model for inference\n",
    "# Place holder for the loaded model\n",
    "loaded_model = Net()\n",
    "loaded_optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)\n",
    "\n",
    "tiledb_model = PyTorchTileDBModel(uri=os.path.basename(tiledb_cloud_model_uri), namespace=namespace, ctx=ctx)\n",
    "tiledb_model.load(model=loaded_model, optimizer=loaded_optimizer)\n",
    "\n",
    "# Check model parameters\n",
    "assert str(model.state_dict()) == str(loaded_model.state_dict())\n",
    "\n",
    "# Check optimizer parameters\n",
    "assert str(optimizer.state_dict()) == str(loaded_optimizer.state_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deregister model and physically delete the array from the S3 bucket\n",
    "tiledb.cloud.array.delete_array(tiledb_cloud_model_uri)"
   ]
  }
 ],
 "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
}