{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5342505d-4a76-404a-9011-7b6af316c6d3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffe8a279-4e33-4f65-86aa-d50114df871f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "TRAIN_EPOCHS = 10 # Number of epochs to train\n",
    "OUTPUTS_MIDDLE_LAYER = 1024 # Number of units for Dense layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3c6c972-a2a4-4ade-8374-e76fc7f03431",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def mnist_datasets():\n",
    "    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(path='/mnt/shared/aie-tutorials/current-release/Data-Science/Kubeflow-GPU/mnist.npz')\n",
    "    x_train, x_test = x_train / np.float32(255), x_test / np.float32(255)\n",
    "    y_train, y_test = y_train.astype(np.int64), y_test.astype(np.int64)\n",
    "    return x_train, x_test, y_train, y_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "794d7272-d629-4a1f-8e17-941be1b5ca37",
   "metadata": {
    "tags": []
   },
   "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(OUTPUTS_MIDDLE_LAYER, activation='relu'),\n",
    "        tf.keras.layers.Dense(10, activation='softmax'),\n",
    "    ])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d4aa20c-f992-4d04-8d88-aef8529730fa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def compile_model(model):\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss='sparse_categorical_crossentropy',\n",
    "                  metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eea5e48c-c1ce-42d5-bc02-d106ea9f6abe",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_model(model, x_train, x_test, y_test, y_train):\n",
    "    # model fitting\n",
    "    start = time.time()\n",
    "    history = model.fit(\n",
    "        x_train, y_train, epochs=TRAIN_EPOCHS,\n",
    "        validation_data=(x_test, y_test)\n",
    "    )\n",
    "    duration_total = time.time() - start\n",
    "    print()\n",
    "    print('Total time %f sec' % (duration_total))\n",
    "    return history, duration_total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3de1a72b-c1a3-40cc-b9e6-2bab67a9d732",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Clean previously saved models\n",
    "model_dir = '/mnt/user/mnist-gpu-test'\n",
    "if tf.io.gfile.exists(model_dir):\n",
    "    tf.io.gfile.rmtree(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b238b1f-8ec6-4bf5-a53b-3d51cb3429db",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load dataset\n",
    "x_train, x_test, y_train, y_test = mnist_datasets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0668c17c-7b36-4b46-87c1-aed208585494",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Train and evaluate for a set number of epochs.\n",
    "with tf.device('/CPU:0'):\n",
    "    # Create the model\n",
    "    model_mnist_cpu = create_model()\n",
    "    # Compile and train\n",
    "    compile_model(model_mnist_cpu)\n",
    "    history_cpu, cpu_duration = train_model(model_mnist_cpu, x_train, x_test, y_test, y_train)\n",
    "    cpu_acc = history_cpu.history['accuracy'][-1]\n",
    "    print(f\"Accuracy: {cpu_acc * 100}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53c5a324-7b2c-4660-89da-b6406c750c7b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Train and evaluate for a set number of epochs.\n",
    "with tf.device('/GPU:0'):\n",
    "    print(\"Num GPUs Available: \", len(tf.config.list_physical_devices('GPU')))\n",
    "    print()\n",
    "    # Create the model\n",
    "    model_mnist_gpu = create_model()\n",
    "    # Compile and train\n",
    "    compile_model(model_mnist_gpu)\n",
    "    history_gpu, gpu_duration = train_model(model_mnist_gpu, x_train, x_test, y_test, y_train)\n",
    "    gpu_acc = history_gpu.history['accuracy'][-1]\n",
    "    print(f\"Accuracy: {gpu_acc * 100}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adb4b3c2-e507-460e-a560-a1f164ca70e1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Save model\n",
    "export_path_cpu = os.path.join(model_dir, 'cpu', 'export')\n",
    "export_path_gpu = os.path.join(model_dir, 'gpu', 'export')\n",
    "tf.saved_model.save(model_mnist_cpu, export_path_cpu)\n",
    "tf.saved_model.save(model_mnist_gpu, export_path_gpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6d872c0-b505-4d6d-8919-e2b678a7f2b7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_accuracy = max(cpu_acc, gpu_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7c85037-0c91-4d89-aaca-b184bfa787d0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"======= Timing =======\")\n",
    "print(f\"CPU time: {cpu_duration} sec\")\n",
    "print(f\"GPU time: {gpu_duration} sec\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81bb0660-6b38-4551-8fe3-f0b69e9fa197",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(model_accuracy)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "kubeflow_notebook": {
   "autosnapshot": false,
   "docker_image": "",
   "experiment": {
    "id": "new",
    "name": "test-gpu-card"
   },
   "experiment_name": "test-gpu-card",
   "katib_metadata": {
    "algorithm": {
     "algorithmName": "grid"
    },
    "maxFailedTrialCount": 3,
    "maxTrialCount": 3,
    "objective": {
     "additionalMetricNames": [],
     "objectiveMetricName": "model-accuracy",
     "type": "maximize"
    },
    "parallelTrialCount": 1,
    "parameters": [
     {
      "feasibleSpace": {
       "list": [
        "32",
        "512",
        "1024"
       ]
      },
      "name": "OUTPUTS_MIDDLE_LAYER",
      "parameterType": "categorical"
     }
    ]
   },
   "katib_run": false,
   "pipeline_description": "test gpu card",
   "pipeline_name": "mnist-gpu",
   "snapshot_volumes": false,
   "steps_defaults": [
    "label:access-ml-pipeline:true",
    "label:add-ldapcert-secret:true",
    "label:add-sssd-secret:true",
    "label:add-user-s3-secret:true"
   ],
   "volume_access_mode": "rwm",
   "volumes": [
    {
     "annotations": [],
     "mount_point": "/mnt/shared",
     "name": "kubeflow-shared-pvc",
     "size": 1,
     "size_type": "Gi",
     "snapshot": false,
     "snapshot_name": "",
     "type": "pvc"
    },
    {
     "annotations": [],
     "mount_point": "/mnt/user",
     "name": "user-pvc",
     "size": 1,
     "size_type": "Gi",
     "snapshot": false,
     "snapshot_name": "",
     "type": "pvc"
    }
   ]
  },
  "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}