{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TensorFlow 2.0 Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NOTE: This notebook is meant to teach you the latest version of Tensorflow 2.0. Most examples on the web today are still in 1.x, so be careful not to confuse the two when looking up documentation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Tensorflow 2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you would like to run this notebook on your local machine, please make sure you follow steps to install the dependencies. You will need TensorFlow version 2.0, and here are some tips to get you started. Please note that Tensorflow 2.0 is still not in a fully 100% stable release, but it's still usable and more intuitive than TF 1.x.\n",
    "\n",
    "     1. Have the latest version of Anaconda installed on your machine.\n",
    "     2. Create a new conda environment starting from Python 3.7. In this setup example, we'll call it tf_20_env.\n",
    "     3. Run the command: source activate tf_20_env\n",
    "     4. Then pip install TF 2.0 as described here: https://www.tensorflow.org/install/pip\n",
    "     \n",
    "\n",
    "A guide on creating Anaconda enviornments: https://uoa-eresearch.github.io/eresearch-cookbook/recipe/2014/11/20/conda/\n",
    "\n",
    "This will give you an new enviornemnt to play in TF 2.0. Generally, if you plan to also use TensorFlow in your other projects, you might also want to keep a seperate Conda environment or virtualenv in Python 3.7 that has Tensorflow 1.9, so you can switch back and forth at will.\n",
    "     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function, unicode_literals\n",
    "import os\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "from tensorflow.keras.layers import Dense, Flatten, Conv2D\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import Model\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this part, we will work with the MNIST data. We will demonstrate how the built-in dataset works in TensorFlow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist = tf.keras.datasets.mnist\n",
    "\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "x_train, x_test = x_train / 255.0, x_test / 255.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for image in x_train[:3]:\n",
    "    plt.figure()\n",
    "    plt.imshow(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Starting Simple: Using the Keras Sequential API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you will quickly see, there are a ton of ways in TensorFlow to build models and train them. TensorFlow has many APIs that you can use, and all these APIs can be confusing. We start this tutorial with the simplest way to build a model in TensorFlow 2.0: using the Keras Sequential API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will first build the tf.keras.Sequential model by stacking layers.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = tf.keras.models.Sequential([\n",
    "  tf.keras.layers.Flatten(input_shape=(28, 28)),\n",
    "  tf.keras.layers.Dense(128, activation='relu'),\n",
    "  tf.keras.layers.Dropout(0.2),\n",
    "  tf.keras.layers.Dense(10, activation='softmax')\n",
    "])\n",
    "\n",
    "model.compile(optimizer='adam',\n",
    "              loss='sparse_categorical_crossentropy',\n",
    "              metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(x_train, y_train, epochs=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss, test_acc = model.evaluate(x_test, y_test)\n",
    "\n",
    "print('\\nTest accuracy:', test_acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## More advanced: Writing a custom training loop with Gradient Tape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we transition to a more advanced usage of TensorFlow. This API will probably be the one you will use the most, as it provides for more flexibility. Note that you can specify more unique architectures with the forward pass. We will use the same MNIST dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reload the dataset \n",
    "def convert_types(image, label):\n",
    "  image = tf.cast(image, tf.float32)\n",
    "  image /= 255\n",
    "  return image, label\n",
    "\n",
    "dataset, info = tfds.load('mnist', data_dir='gs://tfds-data/datasets', with_info=True, as_supervised=True)\n",
    "mnist_train, mnist_test = dataset['train'], dataset['test']\n",
    "mnist_train = mnist_train.map(convert_types).shuffle(10000).batch(32)\n",
    "mnist_test = mnist_test.map(convert_types).batch(32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we will go ahead and build the model using tf.keras. Note how building the model this way allows for a greater degree of control over the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModel(Model):\n",
    "  def __init__(self):\n",
    "    super(MyModel, self).__init__()\n",
    "    self.conv1 = Conv2D(32, 3, activation='relu')\n",
    "    self.flatten = Flatten()\n",
    "    self.d1 = Dense(128, activation='relu')\n",
    "    self.d2 = Dense(10, activation='softmax')\n",
    "\n",
    "  def call(self, x):\n",
    "    x = self.conv1(x)\n",
    "    x = self.flatten(x)\n",
    "    x = self.d1(x)\n",
    "    return self.d2(x)\n",
    "\n",
    "model = MyModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to specify an optimizer and loss function, along with metrics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_object = tf.keras.losses.SparseCategoricalCrossentropy()\n",
    "optimizer = tf.keras.optimizers.Adam()\n",
    "\n",
    "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
    "train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')\n",
    "test_loss = tf.keras.metrics.Mean(name='test_loss')\n",
    "test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can write our train function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def train_step(image, label):\n",
    "  with tf.GradientTape() as tape:\n",
    "    predictions = model(image)\n",
    "    loss = loss_object(label, predictions)\n",
    "  gradients = tape.gradient(loss, model.trainable_variables)\n",
    "  optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "\n",
    "  train_loss(loss)\n",
    "  train_accuracy(label, predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the tf.function annotation on the function above. This means that this function will be compiled into a graph in the backend, allowing it to run efficiently as TensorFlow can optimize the function for you. This automatic conversion of Python code to its graph representation is called AutoGraph, and this creates callable graphs from Python functions.\n",
    "\n",
    "What is this gradient tape?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.Variable(3.0)\n",
    "with tf.GradientTape() as g:\n",
    "  y = x * x\n",
    "dy_dx = g.gradient(y, x) # Will compute to 6.0\n",
    "print(dy_dx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Trainable variables are automatically watched within the GradientTape (when trainable=True, as we will see this is not always the case later in the notebook). Variables are trainable when created by tf.Variable, but not when created via tf.constant. In those cases, if you want to compute the gradients of non-trainable variables, you can manually watch the tensors using the watch function, as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.constant(3.0)\n",
    "with tf.GradientTape() as g:\n",
    "  g.watch(x)\n",
    "  y = x * x\n",
    "dy_dx = g.gradient(y, x) # Will compute to 6.0\n",
    "print(dy_dx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can define our test function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def test_step(image, label):\n",
    "  predictions = model(image)\n",
    "  t_loss = loss_object(label, predictions)\n",
    "\n",
    "  test_loss(t_loss)\n",
    "  test_accuracy(label, predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are ready to train the model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 3\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "  for image, label in mnist_train:\n",
    "    train_step(image, label)\n",
    "\n",
    "  for test_image, test_label in mnist_test:\n",
    "    test_step(test_image, test_label)\n",
    "\n",
    "  template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}'\n",
    "  print (template.format(epoch+1,\n",
    "                         train_loss.result(),\n",
    "                         train_accuracy.result()*100,\n",
    "                         test_loss.result(),\n",
    "                         test_accuracy.result()*100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save and Load Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()\n",
    "\n",
    "train_labels = train_labels[:1000]\n",
    "test_labels = test_labels[:1000]\n",
    "\n",
    "train_images = train_images[:1000].reshape(-1, 28 * 28) / 255.0\n",
    "test_images = test_images[:1000].reshape(-1, 28 * 28) / 255.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Returns a short sequential model\n",
    "def create_model():\n",
    "  model = tf.keras.models.Sequential([\n",
    "    keras.layers.Dense(512, activation='relu', input_shape=(784,)),\n",
    "    keras.layers.Dropout(0.2),\n",
    "    keras.layers.Dense(10, activation='softmax')\n",
    "  ])\n",
    "\n",
    "  model.compile(optimizer='adam',\n",
    "                loss='sparse_categorical_crossentropy',\n",
    "                metrics=['accuracy'])\n",
    "\n",
    "  return model\n",
    "\n",
    "\n",
    "# Create a basic model instance\n",
    "model = create_model()\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_path = \"training_1/cp.ckpt\"\n",
    "checkpoint_dir = os.path.dirname(checkpoint_path)\n",
    "\n",
    "# Create checkpoint callback\n",
    "cp_callback = tf.keras.callbacks.ModelCheckpoint(checkpoint_path,\n",
    "                                                 save_weights_only=True,\n",
    "                                                 verbose=1)\n",
    "\n",
    "model = create_model()\n",
    "\n",
    "model.fit(train_images, train_labels,  epochs = 10,\n",
    "          validation_data = (test_images,test_labels),\n",
    "          callbacks = [cp_callback])  # pass callback to training\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model()\n",
    "\n",
    "loss, acc = model.evaluate(test_images, test_labels)\n",
    "print(\"Untrained model, accuracy: {:5.2f}%\".format(100*acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights(checkpoint_path)\n",
    "loss,acc = model.evaluate(test_images, test_labels)\n",
    "print(\"Restored model, accuracy: {:5.2f}%\".format(100*acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Saving the entire model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model()\n",
    "\n",
    "model.fit(train_images, train_labels, epochs=5)\n",
    "\n",
    "# Save entire model to a HDF5 file\n",
    "model.save('my_model.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Recreate the exact same model, including weights and optimizer.\n",
    "new_model = keras.models.load_model('my_model.h5')\n",
    "new_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss, acc = new_model.evaluate(test_images, test_labels)\n",
    "print(\"Restored model, accuracy: {:5.2f}%\".format(100*acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Putting it all Together: Transfer Learning Using Pretrained Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we will train a model to distinguish between cats and dogs. We will leverage pretrained models for this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We begin my loading in the dataset, and quickly visualizing it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SPLIT_WEIGHTS = (8, 1, 1)\n",
    "splits = tfds.Split.TRAIN.subsplit(weighted=SPLIT_WEIGHTS)\n",
    "\n",
    "(raw_train, raw_validation, raw_test), metadata = tfds.load(\n",
    "    'cats_vs_dogs', split=list(splits),\n",
    "    with_info=True, as_supervised=True)\n",
    "\n",
    "get_label_name = metadata.features['label'].int2str\n",
    "\n",
    "for image, label in raw_train.take(2):\n",
    "  plt.figure()\n",
    "  plt.imshow(image)\n",
    "  plt.title(get_label_name(label))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will preprocess and shuffle the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMG_SIZE = 128 # All images will be resized to 160x160\n",
    "\n",
    "def format_example(image, label):\n",
    "  image = tf.cast(image, tf.float32)\n",
    "  image = (image/127.5) - 1\n",
    "  image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE))\n",
    "  return image, label\n",
    "\n",
    "train = raw_train.map(format_example)\n",
    "validation = raw_validation.map(format_example)\n",
    "test = raw_test.map(format_example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "SHUFFLE_BUFFER_SIZE = 1000\n",
    "train_batches = train.shuffle(SHUFFLE_BUFFER_SIZE).batch(BATCH_SIZE)\n",
    "validation_batches = validation.batch(BATCH_SIZE)\n",
    "test_batches = test.batch(BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have taken care of our data, we will obtain a pretrained model. For this, we will use the MobileNet model (https://arxiv.org/abs/1801.04381). Note that we are using both its architecture, as well as pretrain weights for the model, trained on imagenet. This architecture will act as a feature extractor for us, and that is why we call it \"base_model\". We will then feed the features it extracts into more layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMG_SHAPE = (IMG_SIZE, IMG_SIZE, 3)\n",
    "\n",
    "# Create the base model from the pre-trained model MobileNet V2\n",
    "base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SHAPE,\n",
    "                                               include_top=False,\n",
    "                                               weights='imagenet')\n",
    "\n",
    "base_model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just to get an idea of this pretrained model, we will pass in a batch of images through it, and examine our output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for image_batch, label_batch in train_batches.take(1):\n",
    "    feature_batch = base_model(image_batch)\n",
    "    print(feature_batch.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the input was (16, 128, 128, 3), with batch size as the first dimension. Thus, this feature extractor we have defined converts each 128x128x3 image to a 5x5x128 feature block. Since we only want to use the pretrained model as a feature extractor, we can go ahead and freeze the base. This means that we do not modify the weights of it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_model.trainable = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we add in the last part of the model, which we will train. First, we will add some average pooling to convert features to a single 1280-element vector per image. After we have that, we will add a single dense layer that outputs the classifications of the images. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "global_average_layer = tf.keras.layers.GlobalAveragePooling2D()\n",
    "feature_batch_average = global_average_layer(feature_batch)\n",
    "print(feature_batch_average.shape)\n",
    "\n",
    "\n",
    "prediction_layer = keras.layers.Dense(1)\n",
    "prediction_batch = prediction_layer(feature_batch_average)\n",
    "print(prediction_batch.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = tf.keras.Sequential([\n",
    "  base_model,\n",
    "  global_average_layer,\n",
    "  prediction_layer\n",
    "])\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_learning_rate = 0.0001\n",
    "model.compile(optimizer=tf.keras.optimizers.RMSprop(lr=base_learning_rate),\n",
    "              loss='binary_crossentropy',\n",
    "              metrics=['accuracy'])\n",
    "num_train, num_val, num_test = (\n",
    "  metadata.splits['train'].num_examples*weight/10\n",
    "  for weight in SPLIT_WEIGHTS\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_epochs = 10\n",
    "steps_per_epoch = round(num_train)//BATCH_SIZE\n",
    "validation_steps = 20\n",
    "\n",
    "loss0,accuracy0 = model.evaluate(validation_batches, steps = validation_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit(train_batches,\n",
    "                    epochs=initial_epochs,\n",
    "                    validation_data=validation_batches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}