sdk tests with papermill

Signed-off-by: Yehudit Kerido <[email protected]>

Author: Yehudit Kerido

examples/v1beta1/sdk/tune-train-from-func.ipynb

@@ -81,17 +81,124 @@
   },
   {
    "cell_type": "markdown",
-   "source": [
-    "## Create Train Script for CNN Model\n",
-    "\n",
-    "This is simple **Convolutional Neural Network (CNN)** model for recognizing hand-written digits using [MNIST Dataset](http://yann.lecun.com/exdb/mnist/).  "
-   ],
+   "id": "ee4a3254",
    "metadata": {
     "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
     "pycharm": {
      "name": "#%% md\n"
     }
-   }
+   },
+   "source": [
+    "## Create Train Script for CNN Model\n",
+    "\n",
+    "This is simple **Convolutional Neural Network (CNN)** model for recognizing hand-written digits using [MNIST Dataset](http://yann.lecun.com/exdb/mnist/).  "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fce87ff7-bd14-40de-aaec-824a80021705",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def train_mnist_model(parameters):\n",
+    "    import tensorflow as tf\n",
+    "    import numpy as np\n",
+    "    import logging\n",
+    "\n",
+    "    logging.basicConfig(\n",
+    "        format=\"%(asctime)s %(levelname)-8s %(message)s\",\n",
+    "        datefmt=\"%Y-%m-%dT%H:%M:%SZ\",\n",
+    "        level=logging.INFO,\n",
+    "    )\n",
+    "    logging.info(\"--------------------------------------------------------------------------------------\")\n",
+    "    logging.info(f\"Input Parameters: {parameters}\")\n",
+    "    logging.info(\"--------------------------------------------------------------------------------------\\n\\n\")\n",
+    "\n",
+    "\n",
+    "    # Get HyperParameters from the input params dict.\n",
+    "    lr = float(parameters[\"lr\"])\n",
+    "    num_epoch = int(parameters[\"num_epoch\"])\n",
+    "\n",
+    "    # Set dist parameters and strategy.\n",
+    "    is_dist = parameters[\"is_dist\"]\n",
+    "    num_workers = parameters[\"num_workers\"]\n",
+    "    batch_size_per_worker = 64\n",
+    "    batch_size_global = batch_size_per_worker * num_workers\n",
+    "    strategy = tf.distribute.MultiWorkerMirroredStrategy(\n",
+    "        communication_options=tf.distribute.experimental.CommunicationOptions(\n",
+    "            implementation=tf.distribute.experimental.CollectiveCommunication.RING\n",
+    "        )\n",
+    "    )\n",
+    "\n",
+    "    # Callback class for logging training.\n",
+    "    # Katib parses metrics in this format: <metric-name>=<metric-value>.\n",
+    "    class CustomCallback(tf.keras.callbacks.Callback):\n",
+    "        def on_epoch_end(self, epoch, logs=None):\n",
+    "            logging.info(\n",
+    "                \"Epoch {}/{}. accuracy={:.4f} - loss={:.4f}\".format(\n",
+    "                    epoch+1, num_epoch, logs[\"accuracy\"], logs[\"loss\"]\n",
+    "                )\n",
+    "            )\n",
+    "\n",
+    "    # Prepare MNIST Dataset.\n",
+    "    def mnist_dataset(batch_size):\n",
+    "        (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()\n",
+    "        x_train = x_train / np.float32(255)\n",
+    "        y_train = y_train.astype(np.int64)\n",
+    "        train_dataset = (\n",
+    "            tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
+    "            .shuffle(60000)\n",
+    "            .repeat()\n",
+    "            .batch(batch_size)\n",
+    "        )\n",
+    "        return train_dataset\n",
+    "\n",
+    "    # Build and compile CNN Model.\n",
+    "    def build_and_compile_cnn_model():\n",
+    "        model = tf.keras.Sequential(\n",
+    "            [\n",
+    "                tf.keras.layers.InputLayer(input_shape=(28, 28)),\n",
+    "                tf.keras.layers.Reshape(target_shape=(28, 28, 1)),\n",
+    "                tf.keras.layers.Conv2D(32, 3, activation=\"relu\"),\n",
+    "                tf.keras.layers.Flatten(),\n",
+    "                tf.keras.layers.Dense(128, activation=\"relu\"),\n",
+    "                tf.keras.layers.Dense(10),\n",
+    "            ]\n",
+    "        )\n",
+    "        model.compile(\n",
+    "            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+    "            optimizer=tf.keras.optimizers.SGD(learning_rate=lr),\n",
+    "            metrics=[\"accuracy\"],\n",
+    "        )\n",
+    "        return model\n",
+    "    \n",
+    "    # Download Dataset.\n",
+    "    dataset = mnist_dataset(batch_size_global)\n",
+    "\n",
+    "    # For dist strategy we should build model under scope().\n",
+    "    if is_dist:\n",
+    "        logging.info(\"Running Distributed Training\")\n",
+    "        logging.info(\"--------------------------------------------------------------------------------------\\n\\n\")\n",
+    "        with strategy.scope():\n",
+    "            model = build_and_compile_cnn_model()\n",
+    "    else:\n",
+    "        logging.info(\"Running Single Worker Training\")\n",
+    "        logging.info(\"--------------------------------------------------------------------------------------\\n\\n\")\n",
+    "        model = build_and_compile_cnn_model()\n",
+    "    \n",
+    "    # Start Training.\n",
+    "    model.fit(\n",
+    "        dataset,\n",
+    "        epochs=num_epoch,\n",
+    "        steps_per_epoch=70,\n",
+    "        callbacks=[CustomCallback()],\n",
+    "        verbose=0,\n",
+    "    )"
+   ]
   },
   {
    "cell_type": "markdown",