{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {},
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##### Copyright 2020 The IREE Authors"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Licensed under the Apache License v2.0 with LLVM Exceptions.\n",
        "# See https://llvm.org/LICENSE.txt for license information.\n",
        "# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# ResNet\n",
        "\n",
        "[ResNet](https://arxiv.org/abs/1512.03385) is a deep neural network architecture for image recognition.\n",
        "\n",
        "This notebook\n",
        "\n",
        "* Constructs a [ResNet50](https://www.tensorflow.org/api_docs/python/tf/keras/applications/ResNet50) model using `tf.keras`, with weights pretrained using the [ImageNet](http://www.image-net.org/) dataset\n",
        "* Compiles that model with IREE\n",
        "* Tests TensorFlow and IREE execution of the model on a sample image"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!python -m pip install --upgrade tensorflow"
      ],
      "metadata": {},
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "!python -m pip install --pre iree-base-compiler iree-base-runtime iree-tools-tf -f https://iree.dev/pip-release-links.html"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Imports and common setup\n",
        "\n",
        "from iree import runtime as ireert\n",
        "from iree.tf.support import module_utils\n",
        "\n",
        "import tensorflow as tf\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "# Print version information for future notebook users to reference.\n",
        "!iree-compile --version\n",
        "print(\"TensorFlow version: \", tf.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Construct a pretrained ResNet model with ImageNet weights\n",
        "\n",
        "# Static shape, including batch size (1).\n",
        "# Can be dynamic once dynamic shape support is ready.\n",
        "INPUT_SHAPE = [1, 224, 224, 3]\n",
        "\n",
        "tf_model = tf.keras.applications.resnet50.ResNet50(\n",
        "    weights=\"imagenet\", include_top=True, input_shape=tuple(INPUT_SHAPE[1:]))\n",
        "\n",
        "# Wrap the model in a tf.Module to compile it with IREE.\n",
        "class ResNetModule(tf.Module):\n",
        "\n",
        "  def __init__(self):\n",
        "    super(ResNetModule, self).__init__()\n",
        "    self.model = tf_model\n",
        "\n",
        "  @tf.function(input_signature=[tf.TensorSpec(INPUT_SHAPE, tf.float32)])\n",
        "  def predict(self, x):\n",
        "    return self.model.call(x, training=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@markdown ### Backend Configuration\n",
        "\n",
        "backend_choice = \"iree_llvmcpu (CPU)\" #@param [ \"iree_vmvx (CPU)\", \"iree_llvmcpu (CPU)\", \"iree_vulkan (GPU/SwiftShader)\" ]\n",
        "backend_choice = backend_choice.split(\" \")[0]\n",
        "backend = module_utils.BackendInfo(backend_choice)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Compile ResNet with IREE\n",
        "# This may take a few minutes.\n",
        "iree_module = backend.compile_from_class(ResNetModule, [\"predict\"])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Load a test image of a [labrador](https://commons.wikimedia.org/wiki/File:YellowLabradorLooking_new.jpg)\n",
        "\n",
        "def load_image(path_to_image):\n",
        "  image = tf.io.read_file(path_to_image)\n",
        "  image = tf.image.decode_image(image, channels=3)\n",
        "  image = tf.image.resize(image, (224, 224))\n",
        "  image = image[tf.newaxis, :]\n",
        "  return image\n",
        "\n",
        "content_path = tf.keras.utils.get_file(\n",
        "    'YellowLabradorLooking_new.jpg',\n",
        "    'https://storage.googleapis.com/download.tensorflow.org/example_images/YellowLabradorLooking_new.jpg')\n",
        "content_image = load_image(content_path)\n",
        "\n",
        "print(\"Test image:\")\n",
        "plt.imshow(content_image.numpy().reshape(224, 224, 3) / 255.0)\n",
        "plt.axis(\"off\")\n",
        "plt.tight_layout()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Model pre- and post-processing\n",
        "input_data = tf.keras.applications.resnet50.preprocess_input(content_image)\n",
        "\n",
        "def decode_result(result):\n",
        "  return tf.keras.applications.resnet50.decode_predictions(result, top=3)[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Run TF model\n",
        "\n",
        "print(\"TF prediction:\")\n",
        "tf_result = tf_model.predict(input_data)\n",
        "print(decode_result(tf_result))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Run the model compiled with IREE\n",
        "\n",
        "print(\"IREE prediction:\")\n",
        "iree_result = iree_module.predict(input_data)\n",
        "print(decode_result(iree_result))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}