Add keras lstm fusion codelab.
PiperOrigin-RevId: 306386305 Change-Id: I987a7861ccd6d5fb92719b4475a8bec9325c26ab
This commit is contained in:
Renjie Liu
TensorFlower Gardener
parent
005ad3dda3
commit
242129341e
|
|
@ -0,0 +1,243 @@
|
|||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "Keras LSTM fusion Codelab.ipynb",
|
||||
"provenance": [],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
}
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Z7gEg4DRBwbO",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"\n",
|
||||
"# Overview\n",
|
||||
"This CodeLab demonstrates how to build a fused TFLite LSTM model for MNIST recognition using Keras, and how to convert it to TensorFlow Lite.\n",
|
||||
"\n",
|
||||
"The CodeLab is very similar to the Keras LSTM [CodeLab](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/experimental_new_converter/keras_lstm.ipynb). However, we're creating fused LSTM ops rather than the unfused versoin.\n",
|
||||
"\n",
|
||||
"Also note: We're not trying to build the model to be a real world application, but only demonstrate how to use TensorFlow Lite. You can a build a much better model using CNN models. For a more canonical lstm codelab, please see [here](https://github.com/keras-team/keras/blob/master/examples/imdb_lstm.py)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "1v9muouWCrLA",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"\n",
|
||||
"# Step 0: Prerequisites\n",
|
||||
"It's recommended to try this feature with the newest TensorFlow nightly pip build."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "P0bSI6A5AWaT",
|
||||
"colab_type": "code",
|
||||
"colab": {}
|
||||
},
|
||||
"source": [
|
||||
"!pip install tf-nightly"
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "zeo4IA1xC4O8",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"# Step 1: Build the MNIST LSTM model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "yMtp56hRBvHe",
|
||||
"colab_type": "code",
|
||||
"colab": {}
|
||||
},
|
||||
"source": [
|
||||
"import numpy as np\n",
|
||||
"import tensorflow.compat.v2 as tf\n",
|
||||
"tf.enable_v2_behavior()"
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "IQtjKcMYC_nD",
|
||||
"colab_type": "code",
|
||||
"colab": {}
|
||||
},
|
||||
"source": [
|
||||
"model = tf.keras.models.Sequential([\n",
|
||||
" tf.keras.layers.Input(shape=(28, 28), name='input'),\n",
|
||||
" tf.keras.layers.LSTM(20, time_major=False, return_sequences=True),\n",
|
||||
" tf.keras.layers.Flatten(),\n",
|
||||
" tf.keras.layers.Dense(10, activation=tf.nn.softmax, name='output')\n",
|
||||
"])\n",
|
||||
"model.compile(optimizer='adam',\n",
|
||||
" loss='sparse_categorical_crossentropy',\n",
|
||||
" metrics=['accuracy'])\n",
|
||||
"model.summary()"
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "qS_79wHVDcri",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"# Step 2: Train & Evaluate the model.\n",
|
||||
"We will train the model using MNIST data."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "-yEAraXGDlcQ",
|
||||
"colab_type": "code",
|
||||
"colab": {}
|
||||
},
|
||||
"source": [
|
||||
"# Load MNIST dataset.\n",
|
||||
"(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
|
||||
"x_train, x_test = x_train / 255.0, x_test / 255.0\n",
|
||||
"x_train = x_train.astype(np.float32)\n",
|
||||
"x_test = x_test.astype(np.float32)\n",
|
||||
"\n",
|
||||
"# Change this to True if you want to test the flow rapidly.\n",
|
||||
"# Train with a small dataset and only 1 epoch. The model will work poorly\n",
|
||||
"# but this provides a fast way to test if the conversion works end to end.\n",
|
||||
"_FAST_TRAINING = False\n",
|
||||
"_EPOCHS = 5\n",
|
||||
"if _FAST_TRAINING:\n",
|
||||
" _EPOCHS = 1\n",
|
||||
" _TRAINING_DATA_COUNT = 1000\n",
|
||||
" x_train = x_train[:_TRAINING_DATA_COUNT]\n",
|
||||
" y_train = y_train[:_TRAINING_DATA_COUNT]\n",
|
||||
"\n",
|
||||
"model.fit(x_train, y_train, epochs=_EPOCHS)\n",
|
||||
"model.evaluate(x_test, y_test, verbose=0)"
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "5pGyWlkJDpMQ",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"# Step 3: Convert the Keras model to TensorFlow Lite model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "tB1NZBUHDogR",
|
||||
"colab_type": "code",
|
||||
"colab": {}
|
||||
},
|
||||
"source": [
|
||||
"run_model = tf.function(lambda x: model(x))\n",
|
||||
"# This is important, let's fix the input size.\n",
|
||||
"BATCH_SIZE = 1\n",
|
||||
"STEPS = 28\n",
|
||||
"INPUT_SIZE = 28\n",
|
||||
"concrete_func = run_model.get_concrete_function(\n",
|
||||
" tf.TensorSpec([BATCH_SIZE, STEPS, INPUT_SIZE], model.inputs[0].dtype))\n",
|
||||
"\n",
|
||||
"# model directory.\n",
|
||||
"MODEL_DIR = \"keras_lstm\"\n",
|
||||
"model.save(MODEL_DIR, save_format=\"tf\", signatures=concrete_func)\n",
|
||||
"\n",
|
||||
"converter = tf.lite.TFLiteConverter.from_saved_model(MODEL_DIR)\n",
|
||||
"tflite_model = converter.convert()"
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "INFyl-J3FAOY",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"# Step 4: Check the converted TensorFlow Lite model.\n",
|
||||
"Now load the TensorFlow Lite model and use the TensorFlow Lite python interpreter to verify the results."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0-b0IKK2FGuO",
|
||||
"colab_type": "code",
|
||||
"colab": {}
|
||||
},
|
||||
"source": [
|
||||
"# Run the model with TensorFlow to get expected results.\n",
|
||||
"TEST_CASES = 10\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",
|
||||
"\n",
|
||||
"for i in range(TEST_CASES):\n",
|
||||
" expected = model.predict(x_test[i:i+1])\n",
|
||||
" interpreter.set_tensor(input_details[0][\"index\"], x_test[i:i+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 TF model.\n",
|
||||
" np.testing.assert_almost_equal(expected, result)\n",
|
||||
" print(\"Done. The result of TensorFlow matches the result of TensorFlow Lite.\")\n",
|
||||
"\n",
|
||||
" # Please note: TfLite fused Lstm kernel is stateful, so we need to reset\n",
|
||||
" # the states.\n",
|
||||
" # Clean up internal states.\n",
|
||||
" interpreter.reset_all_variables()"
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Cf6KC9fbFY5f",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"# Step 5: Let's inspect the converted TFLite model.\n",
|
||||
"\n",
|
||||
"Let's check the model, you can see the LSTM will be in it's fused format.\n",
|
||||
"\n",
|
||||
"\n"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
Loading…
Reference in New Issue