site/ko/lite/examples/jax_conversion/overview.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8vD3L4qeREvg"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qLCxmWRyRMZE"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4k5PoHrgJQOU"
      },
      "source": [
        "# TFLite용 Jax 모델 변환\n",
        "\n",
        "## 개요\n",
        "\n",
        "참고: 이 API는 새로운 것이며 pip install tf-nightly를 통해서만 사용할 수 있습니다. TensorFlow 버전 2.7에서 사용할 수 있습니다. 또한 API는 아직 실험적이며 변경될 수 있습니다.\n",
        "\n",
        "이 CodeLab은 Jax를 사용하여 MNIST 인식을 위한 모델을 구축하는 방법과 이를 TensorFlow Lite로 변환하는 방법을 보여줍니다. 이 코드랩은 또한 훈련 후 양자화를 사용하여 Jax 변환 TFLite 모델을 최적화하는 방법을 보여줍니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i8cfOBcjSByO"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://www.tensorflow.org/lite/examples/jax_conversion/overview\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\">TensorFlow.org에서 보기</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ko/lite/examples/jax_conversion/overview.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Google Colab에서 실행</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ko/lite/examples/jax_conversion/overview.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">GitHub에서 소스 보기</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ko/lite/examples/jax_conversion/overview.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">노트북 다운로드</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lq-T8XZMJ-zv"
      },
      "source": [
        "## 전제 조건\n",
        "\n",
        "최신 TensorFlow 야간 pip 빌드에서 이 기능을 사용하는 것이 좋습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EV04hKdrnE4f"
      },
      "outputs": [],
      "source": [
        "!pip install tf-nightly --upgrade\n",
        "!pip install jax --upgrade\n",
        "!pip install jaxlib --upgrade"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QAeY43k9KM55"
      },
      "source": [
        "## 데이터 준비\n",
        "\n",
        "Keras 데이터셋으로 MNIST 데이터를 다운로드하고 전처리합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qSOPSZJn1_Tj"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import functools\n",
        "\n",
        "import time\n",
        "import itertools\n",
        "\n",
        "import numpy.random as npr\n",
        "\n",
        "import jax.numpy as jnp\n",
        "from jax import jit, grad, random\n",
        "from jax.example_libraries import optimizers\n",
        "from jax.example_libraries import stax\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hdJIt3Da2Qn1"
      },
      "outputs": [],
      "source": [
        "def _one_hot(x, k, dtype=np.float32):\n",
        "  \"\"\"Create a one-hot encoding of x of size k.\"\"\"\n",
        "  return np.array(x[:, None] == np.arange(k), dtype)\n",
        "\n",
        "(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()\n",
        "train_images, test_images = train_images / 255.0, test_images / 255.0\n",
        "train_images = train_images.astype(np.float32)\n",
        "test_images = test_images.astype(np.float32)\n",
        "\n",
        "train_labels = _one_hot(train_labels, 10)\n",
        "test_labels = _one_hot(test_labels, 10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0eFhx85YKlEY"
      },
      "source": [
        "## Jax로 MNIST 모델 빌드"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mi3TKB9nnQdK"
      },
      "outputs": [],
      "source": [
        "def loss(params, batch):\n",
        "  inputs, targets = batch\n",
        "  preds = predict(params, inputs)\n",
        "  return -jnp.mean(jnp.sum(preds * targets, axis=1))\n",
        "\n",
        "def accuracy(params, batch):\n",
        "  inputs, targets = batch\n",
        "  target_class = jnp.argmax(targets, axis=1)\n",
        "  predicted_class = jnp.argmax(predict(params, inputs), axis=1)\n",
        "  return jnp.mean(predicted_class == target_class)\n",
        "\n",
        "init_random_params, predict = stax.serial(\n",
        "    stax.Flatten,\n",
        "    stax.Dense(1024), stax.Relu,\n",
        "    stax.Dense(1024), stax.Relu,\n",
        "    stax.Dense(10), stax.LogSoftmax)\n",
        "\n",
        "rng = random.PRNGKey(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bRtnOBdJLd63"
      },
      "source": [
        "## 모델 학습 및 평가"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SWbYRyj7LYZt"
      },
      "outputs": [],
      "source": [
        "step_size = 0.001\n",
        "num_epochs = 10\n",
        "batch_size = 128\n",
        "momentum_mass = 0.9\n",
        "\n",
        "\n",
        "num_train = train_images.shape[0]\n",
        "num_complete_batches, leftover = divmod(num_train, batch_size)\n",
        "num_batches = num_complete_batches + bool(leftover)\n",
        "\n",
        "def data_stream():\n",
        "  rng = npr.RandomState(0)\n",
        "  while True:\n",
        "    perm = rng.permutation(num_train)\n",
        "    for i in range(num_batches):\n",
        "      batch_idx = perm[i * batch_size:(i + 1) * batch_size]\n",
        "      yield train_images[batch_idx], train_labels[batch_idx]\n",
        "batches = data_stream()\n",
        "\n",
        "opt_init, opt_update, get_params = optimizers.momentum(step_size, mass=momentum_mass)\n",
        "\n",
        "@jit\n",
        "def update(i, opt_state, batch):\n",
        "  params = get_params(opt_state)\n",
        "  return opt_update(i, grad(loss)(params, batch), opt_state)\n",
        "\n",
        "_, init_params = init_random_params(rng, (-1, 28 * 28))\n",
        "opt_state = opt_init(init_params)\n",
        "itercount = itertools.count()\n",
        "\n",
        "print(\"\\nStarting training...\")\n",
        "for epoch in range(num_epochs):\n",
        "  start_time = time.time()\n",
        "  for _ in range(num_batches):\n",
        "    opt_state = update(next(itercount), opt_state, next(batches))\n",
        "  epoch_time = time.time() - start_time\n",
        "\n",
        "  params = get_params(opt_state)\n",
        "  train_acc = accuracy(params, (train_images, train_labels))\n",
        "  test_acc = accuracy(params, (test_images, test_labels))\n",
        "  print(\"Epoch {} in {:0.2f} sec\".format(epoch, epoch_time))\n",
        "  print(\"Training set accuracy {}\".format(train_acc))\n",
        "  print(\"Test set accuracy {}\".format(test_acc))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Y1OZBhfQhOj"
      },
      "source": [
        "## TFLite 모델로 변환합니다.\n",
        "\n",
        "참고로 우리는\n",
        "\n",
        "1. params를 `functools.partial` `predict` func에 인라인합니다.\n",
        "2. `jnp.zeros` 빌드합니다. 이것은 Jax가 모델을 추적하는 데 사용되는 \"자리 표시자\" 텐서입니다.\n",
        "3. `experimental_from_jax` 호출합니다.\n",
        "\n",
        "> - `serving_func` 는 목록으로 래핑됩니다.\n",
        "> - 입력은 지정된 이름과 연결되고 목록에 래핑된 배열로 전달됩니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6pcqKZqdNTmn"
      },
      "outputs": [],
      "source": [
        "serving_func = functools.partial(predict, params)\n",
        "x_input = jnp.zeros((1, 28, 28))\n",
        "converter = tf.lite.TFLiteConverter.experimental_from_jax(\n",
        "    [serving_func], [[('input1', x_input)]])\n",
        "tflite_model = converter.convert()\n",
        "with open('jax_mnist.tflite', 'wb') as f:\n",
        "  f.write(tflite_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sqEhzaJPSPS1"
      },
      "source": [
        "## 변환된 TFLite 모델 확인\n",
        "\n",
        "변환된 모델의 결과를 Jax 모델과 비교하십시오."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "acj2AYzjSlaY"
      },
      "outputs": [],
      "source": [
        "expected = serving_func(train_images[0:1])\n",
        "\n",
        "# Run the model with TensorFlow Lite\n",
        "interpreter = tf.lite.Interpreter(model_content=tflite_model)\n",
        "interpreter.allocate_tensors()\n",
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "interpreter.set_tensor(input_details[0][\"index\"], train_images[0:1, :, :])\n",
        "interpreter.invoke()\n",
        "result = interpreter.get_tensor(output_details[0][\"index\"])\n",
        "\n",
        "# Assert if the result of TFLite model is consistent with the JAX model.\n",
        "np.testing.assert_almost_equal(expected, result, 1e-5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qy9Gp4H2SjBL"
      },
      "source": [
        "## 모델 최적화\n",
        "\n",
        "우리는 제공 할 것입니다 `representative_dataset` 모델을 최적화하기 위해 훈련 후 quantiztion을 할 수 있습니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KI0rLV-Meg-2"
      },
      "outputs": [],
      "source": [
        "def representative_dataset():\n",
        "  for i in range(1000):\n",
        "    x = train_images[i:i+1]\n",
        "    yield [x]\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.experimental_from_jax(\n",
        "    [serving_func], [[('x', x_input)]])\n",
        "tflite_model = converter.convert()\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "converter.representative_dataset = representative_dataset\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
        "tflite_quant_model = converter.convert()\n",
        "with open('jax_mnist_quant.tflite', 'wb') as f:\n",
        "  f.write(tflite_quant_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15xQR3JZS8TV"
      },
      "source": [
        "## 최적화된 모델 평가"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X3oOm0OaevD6"
      },
      "outputs": [],
      "source": [
        "expected = serving_func(train_images[0:1])\n",
        "\n",
        "# Run the model with TensorFlow Lite\n",
        "interpreter = tf.lite.Interpreter(model_content=tflite_quant_model)\n",
        "interpreter.allocate_tensors()\n",
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "interpreter.set_tensor(input_details[0][\"index\"], train_images[0:1, :, :])\n",
        "interpreter.invoke()\n",
        "result = interpreter.get_tensor(output_details[0][\"index\"])\n",
        "\n",
        "# Assert if the result of TFLite model is consistent with the Jax model.\n",
        "np.testing.assert_almost_equal(expected, result, 1e-5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QqHXCNa3myor"
      },
      "source": [
        "## 양자화된 모델 크기 비교\n",
        "\n",
        "양자화된 모델이 원래 모델보다 4배 더 작은 것을 볼 수 있어야 합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "imFPw007juVG"
      },
      "outputs": [],
      "source": [
        "!du -h jax_mnist.tflite\n",
        "!du -h jax_mnist_quant.tflite"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "overview.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}