clearml/examples/frameworks/keras/Allegro_Trains_keras_TB_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "wFJPLbY7w7Vj"
   },
   "source": [
    "# Allegro Trains Keras with TensorBoard example\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/allegroai/trains/blob/master/examples/frameworks/keras/Allegro_Trains_keras_TB_example.ipynb)\n",
    "\n",
    "This example introduces Trains with Keras and TensorBoard functionality, including automatic logging, models, and TensorBoard outputs. You can find more frameworks examples [here](https://github.com/allegroai/trains/tree/master/examples/frameworks).\n",
    "\n",
    "Note: This example is based on the Keras `mnist_mlp.py` example.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "K7HA0KcX3XBf"
   },
   "outputs": [],
   "source": [
    "!pip install trains\n",
    "!pip install tensorflow>=2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZM0EIh-GqZuu"
   },
   "source": [
    "### Create a new task.\n",
    "To create a new Task object, call the `Task.init` method providing it with `project_name` (the project name for the experiment) and `task_name` (the name of the experiment). When `Task.init` executes, a link to the Web UI Results page for the newly generated Task will be printed, and the Task will be updated in real time in the Trains demo server.\n",
    "\n",
    "You can read about the `Task` class in the docs [here](https://allegro.ai/docs/task.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RYXhcm58uVGL"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import utils as np_utils\n",
    "from trains import Task\n",
    "\n",
    "# Start a new task\n",
    "task = Task.init(project_name=\"Colab notebooks\", task_name=\"Keras with TensorBoard example\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6A36rDJ7s5Pb"
   },
   "outputs": [],
   "source": [
    "# Train a simple deep NN on the MNIST dataset.\n",
    "# Gets to 98.40% test accuracy after 20 epochs\n",
    "# (there is *a lot* of margin for parameter tuning).\n",
    "# 2 seconds per epoch on a K520 GPU.\n",
    "\n",
    "\n",
    "class TensorBoardImage(keras.callbacks.TensorBoard):\n",
    "    @staticmethod\n",
    "    def make_image(tensor):\n",
    "        from PIL import Image\n",
    "        import io\n",
    "        tensor = np.stack((tensor, tensor, tensor), axis=2)\n",
    "        height, width, channels = tensor.shape\n",
    "        image = Image.fromarray(tensor)\n",
    "        output = io.BytesIO()\n",
    "        image.save(output, format='PNG')\n",
    "        image_string = output.getvalue()\n",
    "        output.close()\n",
    "        return tf.Summary.Image(height=height,\n",
    "                                width=width,\n",
    "                                colorspace=channels,\n",
    "                                encoded_image_string=image_string)\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        if logs is None:\n",
    "            logs = {}\n",
    "        super(TensorBoardImage, self).on_epoch_end(epoch, logs)\n",
    "        images = self.validation_data[0]  # 0 - data; 1 - labels\n",
    "        img = (255 * images[0].reshape(28, 28)).astype('uint8')\n",
    "\n",
    "        image = self.make_image(img)\n",
    "        summary = tf.Summary(value=[tf.Summary.Value(tag='image', image=image)])\n",
    "        self.writer.add_summary(summary, epoch)\n",
    "\n",
    "\n",
    "# the data, shuffled and split between train and test sets\n",
    "nb_classes = 10\n",
    "(X_train, y_train), (X_test, y_test) = keras.datasets.mnist.load_data()\n",
    "\n",
    "X_train = X_train.reshape(60000, 784).astype('float32') / 255.\n",
    "X_test = X_test.reshape(10000, 784).astype('float32') / 255.\n",
    "print(X_train.shape[0], 'train samples')\n",
    "print(X_test.shape[0], 'test samples')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gFMDBxwN4nR2"
   },
   "outputs": [],
   "source": [
    "# convert class vectors to binary class matrices\n",
    "Y_train = np_utils.to_categorical(y_train, nb_classes)\n",
    "Y_test = np_utils.to_categorical(y_test, nb_classes)\n",
    "\n",
    "model = keras.models.Sequential()\n",
    "model.add(keras.layers.Dense(512, input_shape=(784,)))\n",
    "model.add(keras.layers.Activation('relu'))\n",
    "\n",
    "model.add(keras.layers.Dense(512))\n",
    "model.add(keras.layers.Activation('relu'))\n",
    "\n",
    "model.add(keras.layers.Dense(10))\n",
    "model.add(keras.layers.Activation('softmax'))\n",
    "\n",
    "model2 = keras.models.Sequential()\n",
    "model2.add(keras.layers.Dense(512, input_shape=(784,)))\n",
    "model2.add(keras.layers.Activation('relu'))\n",
    "\n",
    "model.summary()\n",
    "\n",
    "model.compile(loss='categorical_crossentropy',\n",
    "              optimizer=keras.optimizers.RMSprop(),\n",
    "              metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "40iQp_Wq4K28"
   },
   "outputs": [],
   "source": [
    "# Advanced: setting model class enumeration and set it for the task\n",
    "labels = dict(('digit_%d' % i, i) for i in range(10))\n",
    "task.set_model_label_enumeration(labels)\n",
    "\n",
    "output_folder = os.path.join(tempfile.gettempdir(), 'keras_example')\n",
    "\n",
    "board = keras.callbacks.TensorBoard(histogram_freq=1, log_dir=output_folder, write_images=False)\n",
    "model_store = keras.callbacks.ModelCheckpoint(filepath=os.path.join(output_folder, 'weight.{epoch}.hdf5'))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5FIKDIzy4YF6"
   },
   "outputs": [],
   "source": [
    "# Fit and evaluate the model\n",
    "\n",
    "history = model.fit(X_train,\n",
    "                    Y_train,\n",
    "                    batch_size=128,\n",
    "                    epochs=6,\n",
    "                    callbacks=[board, model_store],\n",
    "                    verbose=1,\n",
    "                    validation_data=(X_test, Y_test))\n",
    "score = model.evaluate(X_test, Y_test, verbose=0)\n",
    "print('Test score:', score[0])\n",
    "print('Test accuracy:', score[1])"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Allegro Trains keras TB example.ipynb",
   "provenance": []
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}