keras-cifar10-original-train-code-with-2-added-lines.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A simple example on training a CNN model with cifar10 data set."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up keras\n",
    "\n",
    "Import TensorFlow into your program to get started:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow import keras\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlrun\n",
    "import mlrun.frameworks.tf_keras as mlrun_tf_keras"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load a dataset\n",
    "\n",
    "Load the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_datasets(batch_size: int=64):\n",
    "\n",
    "    # Load the data:\n",
    "    (train_x, train_y), (test_x, test_y) = keras.datasets.cifar10.load_data()\n",
    "\n",
    "    # Scale the images:\n",
    "    train_x = train_x.astype(np.float32)\n",
    "    train_x /= 255.0\n",
    "    test_x = test_x.astype(np.float32)\n",
    "    test_x /= 255.0\n",
    "\n",
    "    # One hot encode target values\n",
    "    train_y = keras.utils.to_categorical(train_y)\n",
    "    test_y = keras.utils.to_categorical(test_y)\n",
    "\n",
    "    # create data generator\n",
    "    image_data_generator = keras.preprocessing.image.ImageDataGenerator(\n",
    "        width_shift_range=0.1, height_shift_range=0.1, horizontal_flip=True\n",
    "    )\n",
    "\n",
    "    # Wrap up the datasets:\n",
    "    training_set = image_data_generator.flow(train_x, train_y, batch_size=batch_size)\n",
    "    validation_set = (test_x, test_y)\n",
    "\n",
    "    return training_set, validation_set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build a machine learning model\n",
    "\n",
    "Build a `keras.models.Sequential` model by stacking layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model():\n",
    "    # Initialize a sequential model:\n",
    "    model = keras.models.Sequential()\n",
    "\n",
    "    # Add the input block:\n",
    "    model.add(keras.layers.InputLayer(input_shape=(32, 32, 3)))\n",
    "\n",
    "    # Add VGG blocks:\n",
    "    filters_list = [32, 64, 128]\n",
    "    dropout_list = [0.3, 0.5, 0.5]\n",
    "    for filters, dropout in zip(filters_list, dropout_list):\n",
    "        model.add(keras.layers.Conv2D(filters, (3, 3), padding=\"same\", activation=\"relu\"))\n",
    "        model.add(keras.layers.BatchNormalization())\n",
    "        model.add(keras.layers.Conv2D(filters, (3, 3), padding=\"same\", activation=\"relu\"))\n",
    "        model.add(keras.layers.BatchNormalization())\n",
    "        model.add(keras.layers.MaxPooling2D(pool_size=(2, 2)))\n",
    "        model.add(keras.layers.Dropout(dropout))\n",
    "\n",
    "    # Add the output block:\n",
    "    model.add(keras.layers.Flatten())\n",
    "    model.add(keras.layers.Dense(filters_list[-1], activation=\"relu\"))\n",
    "    model.add(keras.layers.BatchNormalization())\n",
    "    model.add(keras.layers.Dropout(dropout_list[-1]))\n",
    "    model.add(keras.layers.Dense(10, activation=\"softmax\"))\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define training function\n",
    "\n",
    "Use the `Model.fit` method to adjust your model parameters and minimize the loss: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(batch_size: int=64, lr: float=0.001, epochs: int = 3):\n",
    "    \n",
    "    print(f'batch size    === >>> {batch_size}')\n",
    "    print(f'learning rate === >>> {lr}')\n",
    "    print(f'epochs        === >>> {epochs}')\n",
    "    \n",
    "    # Get the datasets:\n",
    "    training_set, validation_set = get_datasets(batch_size)\n",
    "\n",
    "    # Get the model:\n",
    "    model = get_model()\n",
    "\n",
    "    ################### Apply MLRun here: ###################\n",
    "    context = mlrun.get_or_create_ctx(name=\"cifar10_ctx\")\n",
    "    mlrun_tf_keras.apply_mlrun(model=model, context=context)\n",
    "    ################### Apply MLRun here: ###################\n",
    "    \n",
    "    # Compile the model:\n",
    "    optimizer = keras.optimizers.Adam(learning_rate=lr)\n",
    "    model.compile(\n",
    "        optimizer=optimizer, loss=\"categorical_crossentropy\", metrics=[\"accuracy\"]\n",
    "    )\n",
    "\n",
    "    # Train the model:\n",
    "    model.fit(training_set, epochs=epochs, validation_data=validation_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now you dno't train your model by running the train() function in this nodebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}