[Enhancement] Support Subclassed Model API and custom objects (#74)

* Saving custom models functionality

Author: ktsitsi

examples/models/tensorflow_keras_tiledb_models_example.ipynb

@@ -526,7 +526,6 @@
    "execution_count": null,
    "metadata": {
     "pycharm": {
-     "is_executing": true,
      "name": "#%%\n"
     }
    },
@@ -535,12 +534,140 @@
     "tiledb.ls('MNIST_Group', lambda obj_path, obj_type: print(obj_path, obj_type))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Model Subclassing\n",
+    "\n",
+    "Apart from being able to store models, which have been created with Symbolic APIs\n",
+    "(Sequential, Functional) someone can store models that are being designed based on\n",
+    "Imperative API (aka. Model Subclassing).\n",
+    "\n",
+    "Let's first design a simple model:"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from tensorflow import keras\n",
+    "\n",
+    "class CustomModel(keras.Model):\n",
+    "    def __init__(self, hidden_units):\n",
+    "        super(CustomModel, self).__init__()\n",
+    "        self.hidden_units = hidden_units\n",
+    "        self.dense_layers = [keras.layers.Dense(u) for u in hidden_units]\n",
+    "\n",
+    "    def call(self, inputs):\n",
+    "        x = inputs\n",
+    "        for layer in self.dense_layers:\n",
+    "            x = layer(x)\n",
+    "        return x\n",
+    "\n",
+    "    def get_config(self):\n",
+    "        return {\"hidden_units\": self.hidden_units}\n",
+    "\n",
+    "    @classmethod\n",
+    "    def from_config(cls, config):\n",
+    "        return cls(**config)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Then we can create a trivial input dataset for testing the model. Remember that\n",
+    "for custom models to be initialised they need to be called on data."
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "model = CustomModel([16, 16, 10])\n",
+    "# Build the model by calling it\n",
+    "input_arr = tf.random.uniform((1, 5))\n",
+    "outputs = model(input_arr)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "We then can save the model as a TileDB array."
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "tiledb_model_custom = TensorflowKerasTileDBModel(uri='tiledb-keras-custom-model', model=model)\n",
+    "tiledb_model_custom.save(include_optimizer=True, update=False)\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Loading the subclassed model requires `custom_objects` to be passed as an argument\n",
+    "and the `input_shape` of the model so it can be built. The output of two models are\n",
+    "exactly the same"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "loaded_custom = tiledb_model_custom.load(custom_objects={\"CustomModel\": CustomModel}, input_shape=(1, 5))\n",
+    "outputs_loaded = loaded_custom((1, 5))\n",
+    "outputs == outputs_loaded"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
   }
  ],
  "metadata": {
@@ -564,4 +691,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 1
-}
+}