chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
License Check / license-check (push) Failing after 2s

This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
+137
View File
@@ -0,0 +1,137 @@
# Copyright 2026 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
book_path: /lite/_book.yaml
project_path: /lite/_project.yaml
title: Models
landing_page:
custom_css_path: /site-assets/css/style.css
nav: left
meta_tags:
- name: description
content: >
Overview of models for TensorFlow Lite
rows:
- classname: devsite-landing-row-100
items:
- description: >
<p>
TensorFlow Lite uses TensorFlow models converted into a smaller, more efficient machine
learning (ML) model format. You can use pre-trained models with TensorFlow Lite, modify
existing models, or build your own TensorFlow models and then convert them to
TensorFlow Lite format. TensorFlow Lite models can perform almost any task a regular
TensorFlow model can do: object detection, natural language processing, pattern
recognition, and more using a wide range of input data including images, video, audio,
and text.
</p>
- classname: devsite-landing-row-100
items:
- description: >
<h2 class="tfo-landing-page-heading no-link">Learning roadmap</h2>
- classname: devsite-landing-row-100
items:
- classname: tfo-landing-page-card
description: >
<a href="/lite/models/convert/index"><h3 class="no-link">Have a TensorFlow model?</h3></a>
Skip to the <a href="/lite/models/convert/index">Convert</a> section for information about
getting your model to run with TensorFlow Lite.
path: /lite/models/convert/index
- classname: tfo-landing-page-card
description: >
<a href="#get_models"><h3 class="no-link">Need a model for TensorFlow Lite?</h3></a>
For guidance on getting models for your use case,
<a href="#get_models">keep reading</a>.
- classname: devsite-landing-row-100
items:
- description: >
<br>
<h2 class="tfo-landing-page-heading no-link" id="get_models">
Get models for TensorFlow Lite</h2>
<p>
You don't have to build a TensorFlow Lite model to start using machine learning on
mobile or edge devices. Many already-built and optimized models are available for you to
use right away in your application. You can start with using pre-trained models in
TensorFlow Lite and move up to building custom models over time, as follows:
</p>
<ol>
<li>Start developing machine learning features with already
<a href="/lite/models/trained">trained models.</a></li>
<li>Modify existing TensorFlow Lite models using tools such as
<a href="/lite/models/modify/model_maker">Model Maker</a>.</li>
<li>Build a
<a href="/tutorials/customization/custom_training_walkthrough">
custom model</a> with TensorFlow tools and then
<a href="/lite/models/convert">convert</a> it to TensorFlow Lite.</li>
</ol>
<br>
<h2 class="tfo-landing-page-heading no-link">Using models for quick tasks: ML Kit</h2>
<p>
If you are trying to quickly implement features or utility tasks with machine
learning, you should review the use cases supported by
<a href="https://developers.google.com/ml-kit">ML Kit</a> before starting
development with TensorFlow Lite. This development tool provides APIs you can call
directly from mobile apps to complete common ML tasks such as barcode scanning and
on-device translation. Using this method can help you get results fast. However, ML Kit
has limited options for extending its capabilities. For more information,
see the <a href="https://developers.google.com/ml-kit">ML Kit</a> developer
documentation.
</p>
<br>
<h2 class="tfo-landing-page-heading no-link">Building models for your app: Constraints
</h2>
<p>
If building a custom model for your specific use case is your ultimate goal, you should
start with developing and training a TensorFlow model or extending an existing one.
Before you start your model development process, you should be aware of the constraints
for TensorFlow Lite models and build your model with these constraints in mind:
<ul>
<li>Limited compute capabilities</li>
<li>Size of models</li>
<li>Size of data</li>
<li>Supported TensorFlow operations</li>
</ul>
</p>
<p>
For more detail about each of these constraints, see
<a href="./build#model_design_constraints">model design contraints</a> in the
Model build overview.
For more information about building effective, compatible, high performance models for
TensorFlow Lite, see <a href="../performance/best_practices">
Performance best practices</a>.
</p>
- classname: devsite-landing-row-100
items:
- description: >
<h2 class="tfo-landing-page-heading no-link">Next steps</h2>
- classname: devsite-landing-row-100
items:
- classname: tfo-landing-page-card
description: >
<a href="/lite/models/trained"><h3 class="no-link">Pick trained model</h3></a>
Learn how to pick a pre-trained ML model to use with TensorFlow Lite.
path: /lite/models/trained
- classname: tfo-landing-page-card
description: >
<a href="/lite/models/modify/model_maker"><h3 class="no-link">Modify models</h3></a>
Use TensorFlow Lite Model Maker to modify models using your training data.
path: /lite/models/modify/model_maker
- classname: tfo-landing-page-card
description: >
<a href="/lite/models/convert/index"><h3 class="no-link">Build models</h3></a>
Learn how to build custom TensorFlow models to use with TensorFlow Lite.
path: /lite/performance/best_practices
@@ -0,0 +1,48 @@
# API Updates <a name="api_updates"></a>
This page provides information about updates made to the
`tf.lite.TFLiteConverter` [Python API](index.md) in TensorFlow 2.x.
Note: If any of the changes raise concerns, please file a
[GitHub issue](https://github.com/tensorflow/tensorflow/issues/new?template=60-tflite-converter-issue.md).
* TensorFlow 2.3
* Support integer (previously, only float) input/output type for integer
quantized models using the new `inference_input_type` and
`inference_output_type` attributes. Refer to this
[example usage](../../performance/post_training_quantization.md#integer_only).
* Support conversion and resizing of models with dynamic dimensions.
* Added a new experimental quantization mode with 16-bit activations and
8-bit weights.
* TensorFlow 2.2
* By default, leverage [MLIR-based conversion](https://mlir.llvm.org/),
Google's cutting edge compiler technology for machine learning. This
enables conversion of new classes of models, including Mask R-CNN,
Mobile BERT, etc and supports models with functional control flow.
* TensorFlow 2.0 vs TensorFlow 1.x
* Renamed the `target_ops` attribute to `target_spec.supported_ops`
* Removed the following attributes:
* _quantization_: `inference_type`, `quantized_input_stats`,
`post_training_quantize`, `default_ranges_stats`,
`reorder_across_fake_quant`, `change_concat_input_ranges`,
`get_input_arrays()`. Instead,
[quantize aware training](https://www.tensorflow.org/model_optimization/guide/quantization/training)
is supported through the `tf.keras` API and
[post training quantization](../../performance/post_training_quantization.md)
uses fewer attributes.
* _visualization_: `output_format`, `dump_graphviz_dir`,
`dump_graphviz_video`. Instead, the recommended approach for
visualizing a TensorFlow Lite model is to use
[visualize.py](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/tools/visualize.py).
* _frozen graphs_: `drop_control_dependency`, as frozen graphs are
unsupported in TensorFlow 2.x.
* Removed other converter APIs such as `tf.lite.toco_convert` and
`tf.lite.TocoConverter`
* Removed other related APIs such as `tf.lite.OpHint` and
`tf.lite.constants` (the `tf.lite.constants.*` types have been mapped to
`tf.*` TensorFlow data types, to reduce duplication)
@@ -0,0 +1,231 @@
# Convert TensorFlow models
This page describes how to convert a TensorFlow model
to a TensorFlow Lite model (an optimized
[FlatBuffer](https://google.github.io/flatbuffers/) format identified by the
`.tflite` file extension) using the TensorFlow Lite converter.
Note: This guide assumes you've both
[installed TensorFlow 2.x](https://www.tensorflow.org/install/pip#tensorflow-2-packages-are-available)
and trained models in TensorFlow 2.x.
If your model is trained in TensorFlow 1.x, considering
[migrating to TensorFlow 2.x](https://www.tensorflow.org/guide/migrate/tflite).
To identify the installed TensorFlow version, run
`print(tf.__version__)`.
## Conversion workflow
The diagram below illustrations the high-level workflow for converting
your model:
![TFLite converter workflow](../../images/convert/convert.png)
**Figure 1.** Converter workflow.
You can convert your model using one of the following options:
1. [Python API](#python_api) (***recommended***):
This allows you to integrate the conversion into your development pipeline,
apply optimizations, add metadata and many other tasks that simplify
the conversion process.
2. [Command line](#cmdline): This only supports basic model conversion.
Note: In case you encounter any issues during model conversion, create a
[GitHub issue](https://github.com/tensorflow/tensorflow/issues/new?template=60-tflite-converter-issue.md).
## Python API <a name="python_api"></a>
*Helper code: To learn more about the TensorFlow Lite converter
API, run `print(help(tf.lite.TFLiteConverter))`.*
Convert a TensorFlow model using
[`tf.lite.TFLiteConverter`](https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter).
A TensorFlow model is stored using the SavedModel format and is
generated either using the high-level `tf.keras.*` APIs (a Keras model) or
the low-level `tf.*` APIs (from which you generate concrete functions). As a
result, you have the following three options (examples are in the next few
sections):
* `tf.lite.TFLiteConverter.from_saved_model()` (**recommended**): Converts
a [SavedModel](https://www.tensorflow.org/guide/saved_model).
* `tf.lite.TFLiteConverter.from_keras_model()`: Converts a
[Keras](https://www.tensorflow.org/guide/keras/overview) model.
* `tf.lite.TFLiteConverter.from_concrete_functions()`: Converts
[concrete functions](https://www.tensorflow.org/guide/intro_to_graphs).
### Convert a SavedModel (recommended) <a name="saved_model"></a>
The following example shows how to convert a
[SavedModel](https://www.tensorflow.org/guide/saved_model) into a TensorFlow
Lite model.
```python
import tensorflow as tf
# Convert the model
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) # path to the SavedModel directory
tflite_model = converter.convert()
# Save the model.
with open('model.tflite', 'wb') as f:
f.write(tflite_model)
```
### Convert a Keras model <a name="keras"></a>
The following example shows how to convert a
[Keras](https://www.tensorflow.org/guide/keras/overview) model into a TensorFlow
Lite model.
```python
import tensorflow as tf
# Create a model using high-level tf.keras.* APIs
model = tf.keras.models.Sequential([
tf.keras.layers.Dense(units=1, input_shape=[1]),
tf.keras.layers.Dense(units=16, activation='relu'),
tf.keras.layers.Dense(units=1)
])
model.compile(optimizer='sgd', loss='mean_squared_error') # compile the model
model.fit(x=[-1, 0, 1], y=[-3, -1, 1], epochs=5) # train the model
# (to generate a SavedModel) tf.saved_model.save(model, "saved_model_keras_dir")
# Convert the model.
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
# Save the model.
with open('model.tflite', 'wb') as f:
f.write(tflite_model)
```
### Convert concrete functions <a name="concrete_function"></a>
The following example shows how to convert
[concrete functions](https://www.tensorflow.org/guide/intro_to_graphs) into a
TensorFlow Lite model.
```python
import tensorflow as tf
# Create a model using low-level tf.* APIs
class Squared(tf.Module):
@tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32)])
def __call__(self, x):
return tf.square(x)
model = Squared()
# (ro run your model) result = Squared(5.0) # This prints "25.0"
# (to generate a SavedModel) tf.saved_model.save(model, "saved_model_tf_dir")
concrete_func = model.__call__.get_concrete_function()
# Convert the model.
converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func],
model)
tflite_model = converter.convert()
# Save the model.
with open('model.tflite', 'wb') as f:
f.write(tflite_model)
```
### Other features
* Apply [optimizations](../../performance/model_optimization.md). A common
optimization used is
[post training quantization](../../performance/post_training_quantization.md),
which can further reduce your model latency and size with minimal loss in
accuracy.
* Add [metadata](metadata.md), which makes it easier to create platform
specific wrapper code when deploying models on devices.
### Conversion errors
The following are common conversion errors and their solutions:
* Error: `Some ops are not supported by the native TFLite runtime, you can
enable TF kernels fallback using TF Select. See instructions:
https://www.tensorflow.org/lite/guide/ops_select. TF Select ops: ..., ..,
...`
Solution: The error occurs as your model has TF ops that don't have a
corresponding TFLite implementation. You can resolve this by
[using the TF op in the TFLite model](../../guide/ops_select.md)
(recommended). If you want to generate a model with TFLite ops only, you can
either add a request for the missing TFLite op in
[Github issue #21526](https://github.com/tensorflow/tensorflow/issues/21526)
(leave a comment if your request hasnt already been mentioned) or
[create the TFLite op](../../guide/ops_custom.md#create_and_register_the_operator)
yourself.
* Error: `.. is neither a custom op nor a flex op`
Solution: If this TF op is:
* Supported in TF: The error occurs because the TF op is missing from the
[allowlist](../../guide/op_select_allowlist.md) (an exhaustive list of
TF ops supported by TFLite). You can resolve this as follows:
1. [Add missing ops to the allowlist](../../guide/op_select_allowlist.md#add_tensorflow_core_operators_to_the_allowed_list).
2. [Convert the TF model to a TFLite model and run inference](../../guide/ops_select.md).
* Unsupported in TF: The error occurs because TFLite is unaware of the
custom TF operator defined by you. You can resolve this as follows:
1. [Create the TF op](https://www.tensorflow.org/guide/create_op).
2. [Convert the TF model to a TFLite model](../../guide/op_select_allowlist.md#users_defined_operators).
3. [Create the TFLite op](../../guide/ops_custom.md#create_and_register_the_operator)
and run inference by linking it to the TFLite runtime.
## Command Line Tool <a name="cmdline"></a>
**Note:** It is highly recommended that you use the [Python API](#python_api)
listed above instead, if possible.
If you've
[installed TensorFlow 2.x from pip](https://www.tensorflow.org/install/pip), use
the `tflite_convert` command. To view all the available flags, use the
following command:
```sh
$ tflite_convert --help
`--output_file`. Type: string. Full path of the output file.
`--saved_model_dir`. Type: string. Full path to the SavedModel directory.
`--keras_model_file`. Type: string. Full path to the Keras H5 model file.
`--enable_v1_converter`. Type: bool. (default False) Enables the converter and flags used in TF 1.x instead of TF 2.x.
You are required to provide the `--output_file` flag and either the `--saved_model_dir` or `--keras_model_file` flag.
```
If you have the
[TensorFlow 2.x source](https://www.tensorflow.org/install/source)
downloaded and want to run the converter from that source without building and
installing the package,
you can replace '`tflite_convert`' with
'`bazel run tensorflow/lite/python:tflite_convert --`' in the command.
### Converting a SavedModel <a name="cmdline_saved_model"></a>
```sh
tflite_convert \
--saved_model_dir=/tmp/mobilenet_saved_model \
--output_file=/tmp/mobilenet.tflite
```
### Converting a Keras H5 model <a name="cmdline_keras_model"></a>
```sh
tflite_convert \
--keras_model_file=/tmp/mobilenet_keras_model.h5 \
--output_file=/tmp/mobilenet.tflite
```
## Next Steps
Use the [TensorFlow Lite interpreter](../../guide/inference.md) to run inference
on a client device (e.g. mobile, embedded).
@@ -0,0 +1,150 @@
# Model conversion overview
The machine learning (ML) models you use with TensorFlow Lite are originally
built and trained using TensorFlow core libraries and tools. Once you've built
a model with TensorFlow core, you can convert it to a smaller, more
efficient ML model format called a TensorFlow Lite model.
This section provides guidance for converting
your TensorFlow models to the TensorFlow Lite model format.
Note: If you don't have a model to convert yet, see the
[Models overview](../)
page for guidance on choosing or building models.
## Conversion workflow
Converting TensorFlow models to TensorFlow Lite format can take a few paths
depending on the content of your ML model. As the first step of that process,
you should evaluate your model to determine if it can be directly converted.
This evaluation determines if the content of the model is supported by the
standard TensorFlow Lite runtime environments based on the TensorFlow operations
it uses. If your model uses operations outside of the supported set, you have
the option to refactor your model or use advanced conversion techniques.
The diagram below shows the high level steps in converting a model.
![TFLite conversion workflow](../../images/convert/convert_workflow_diag.png)
**Figure 1.** TensorFlow Lite conversion workflow.
The following sections outline the process of evaluating and converting models
for use with TensorFlow Lite.
### Input model formats
You can use the converter with the following input model formats:
* [SavedModel](https://www.tensorflow.org/guide/saved_model)
(***recommended***): A TensorFlow model saved as a set of files on disk.
* [Keras model](https://www.tensorflow.org/guide/keras/overview):
A model created using the high level Keras API.
* [Keras H5 format](https://www.tensorflow.org/guide/keras/save_and_serialize#keras_h5_format):
A light-weight alternative to SavedModel format supported by Keras API.
* [Models built from concrete functions](https://www.tensorflow.org/guide/intro_to_graphs):
A model created using the low level TensorFlow API.
You can save both the Keras and concrete function models as a SavedModel
and convert using the recommeded path.
Note: To avoid errors during inference, include signatures when exporting to the
SavedModel format.
The TensorFlow converter supports converting TensorFlow model's
input/output specifications to TensorFlow Lite models. See the topic
on [adding signatures](https://tensorflow.org/lite/guide/signatures).
If you have a Jax model, you can use the `TFLiteConverter.experimental_from_jax`
API to convert it to the TensorFlow Lite format. Note that this API is subject
to change while in experimental mode.
### Conversion evaluation
Evaluating your model is an important step before attempting to convert it.
When evaluating,
you want to determine if the contents of your model is compatible with the
TensorFlow Lite format. You should also determine if your model is a good fit
for use on mobile and edge devices in terms of the size of data the model uses,
its hardware processing requirements, and the model's overall size and
complexity.
For many models, the converter should work out of the box. However,
TensorFlow Lite builtin operator library supports a subset of
TensorFlow core operators, which means some models may need additional
steps before converting to TensorFlow Lite.
Additionally some operations that are supported by TensorFlow Lite have
restricted usage requirements for performance reasons. See the
[operator compatibility](../../guide/ops_compatibility) guide
to determine if your model needs to be refactored for conversion.
Key Point: Most models can be directly converted to TensorFlow Lite format. Some
models may require refactoring or use of advanced conversion techniques to
make them compatible.
### Model conversion
The TensorFlow Lite converter takes a TensorFlow model and generates a
TensorFlow Lite model (an optimized
[FlatBuffer](https://google.github.io/flatbuffers/) format identified by the
`.tflite` file extension). You can load
a SavedModel or directly convert a model you create in code.
The converter takes 3 main flags (or options) that customize the conversion
for your model:
1. [Compatibility flags](../../guide/ops_compatibility) allow you to specify
whether the conversion should allow custom operators.
1. [Optimization flags](../../performance/model_optimization) allow you to
specify the type of optimization to apply
during conversion. The most commonly used optimization technique is
[post-training quantization]().
1. [Metadata flags](metadata) allow you to add metadata to the converted model
which makes it easier to create platform specific wrapper code when deploying
models on devices.
You can convert your model using the [Python API](convert_models#python_api) or
the [Command line](convert_models#cmdline) tool. See the
[Convert TF model](convert_models) guide for step by step
instructions on running the converter on your model.
Typically you would convert your model for the standard TensorFlow Lite
[runtime environment](../../android#runtime) or the
[Google Play services runtime environment](../../android/play_services)
for TensorFlow Lite (Beta). Some advanced use cases require
customization of model runtime environment, which require additional steps in
the conversion proceess. See the
[advanced runtime environment](../../android#adv_runtime) section of the Android
overview for more guidance.
## Advanced conversion
If you run into [errors](convert_models#conversion_errors)
while running the converter on your model, it's most likely that you have an
operator compatibility issue. Not all TensorFlow operations are
supported by TensorFlow
Lite. You can work around these issues by refactoring your model, or by using
advanced conversion options that allow you to create a modified TensorFlow Lite
format model and a custom runtime environment for that model.
* See the [Model compatibility overview](../../guide/ops_compatibility)
for more information on TensorFlow and TensorFlow Lite model compatibility
considerations.
* Topics under the Model compatibility overview cover advanced techniques for
refactoring your model, such as the [Select operators](../../guide/ops_select)
guide.
* For full list of operations and limitations see
[TensorFlow Lite Ops page](https://www.tensorflow.org/mlir/tfl_ops).
## Next steps
* See the [convert TF models](convert_models) guide to quickly get started on
converting your model.
* See the [optimization overview](../../performance/model_optimization) for
guidance on how to optimize your converted model using techniques like
[post-training quantization](../../performance/post_training_quantization).
* See the [Adding metadata overview](metadata) to learn how to add metadata to
your models. Metadata provides other uses a description of your model as well
as information that can be leveraged by code generators.
@@ -0,0 +1,546 @@
# Adding metadata to TensorFlow Lite models
TensorFlow Lite metadata provides a standard for model descriptions. The
metadata is an important source of knowledge about what the model does and its
input / output information. The metadata consists of both
* human readable parts which convey the best practice when using the model,
and
* machine readable parts that can be leveraged by code generators, such as the
[TensorFlow Lite Android code generator](../../inference_with_metadata/codegen.md#generate-model-interfaces-with-tensorflow-lite-code-generator-codegen)
and the
[Android Studio ML Binding feature](../../inference_with_metadata/codegen.md#use-android-studio-ml-model-binding-mlbinding).
All image models published on
[TensorFlow Hub](https://tfhub.dev/s?deployment-format=lite) have been populated
with metadata.
## Model with metadata format
<center><img src="../../images/convert/model_with_metadata.png" alt="model_with_metadata" width="70%"></center>
<center>Figure 1. TFLite model with metadata and associated files.</center>
Model metadata is defined in
[metadata_schema.fbs](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/metadata_schema.fbs),
a
[FlatBuffer](https://google.github.io/flatbuffers/index.html#flatbuffers_overview)
file. As shown in Figure 1, it is stored in the
[metadata](https://github.com/tensorflow/tensorflow/blob/bd73701871af75539dd2f6d7fdba5660a8298caf/tensorflow/lite/schema/schema.fbs#L1208)
field of the
[TFLite model schema](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/schema/schema.fbs),
under the name, `"TFLITE_METADATA"`. Some models may come with associated files,
such as
[classification label files](https://github.com/tensorflow/examples/blob/dd98bc2b595157c03ac9fa47ac8659bb20aa8bbd/lite/examples/image_classification/android/models/src/main/assets/labels.txt#L1).
These files are concatenated to the end of the original model file as a ZIP
using the ZipFile
["append" mode](https://pymotw.com/2/zipfile/#appending-to-files) (`'a'` mode).
TFLite Interpreter can consume the new file format in the same way as before.
See [Pack the associated files](#pack-the-associated-files) for more
information.
See the instruction below about how to populate, visualize, and read metadata.
## Setup the metadata tools
Before adding metadata to your model, you will need to a Python programming
environment setup for running TensorFlow. There is a detailed guide on how to
set this up [here](https://www.tensorflow.org/install).
After setup the Python programming environment, you will need to install
additional tooling:
```sh
pip install tflite-support
```
TensorFlow Lite metadata tooling supports Python 3.
## Adding metadata using Flatbuffers Python API
Note: to create metadata for the popular ML tasks supported in
[TensorFlow Lite Task Library](https://ai.google.dev/edge/litert/libraries/task_library/overview),
use the high-level API in the
[TensorFlow Lite Metadata Writer Library](metadata_writer_tutorial.ipynb).
There are three parts to the model metadata in the
[schema](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/metadata_schema.fbs):
1. **Model information** - Overall description of the model as well as items
such as license terms. See
[ModelMetadata](https://github.com/tensorflow/tflite-support/blob/4cd0551658b6e26030e0ba7fc4d3127152e0d4ae/tensorflow_lite_support/metadata/metadata_schema.fbs#L640).
2. **Input information** - Description of the inputs and pre-processing
required such as normalization. See
[SubGraphMetadata.input_tensor_metadata](https://github.com/tensorflow/tflite-support/blob/4cd0551658b6e26030e0ba7fc4d3127152e0d4ae/tensorflow_lite_support/metadata/metadata_schema.fbs#L590).
3. **Output information** - Description of the output and post-processing
required such as mapping to labels. See
[SubGraphMetadata.output_tensor_metadata](https://github.com/tensorflow/tflite-support/blob/4cd0551658b6e26030e0ba7fc4d3127152e0d4ae/tensorflow_lite_support/metadata/metadata_schema.fbs#L599).
Since TensorFlow Lite only supports single subgraph at this point, the
[TensorFlow Lite code generator](../../inference_with_metadata/codegen..md#generate-model-interfaces-with-tensorflow-lite-code-generator-codegen)
and the
[Android Studio ML Binding feature](../../inference_with_metadata/codegen.md#use-android-studio-ml-model-binding-mlbinding)
will use `ModelMetadata.name` and `ModelMetadata.description`, instead of
`SubGraphMetadata.name` and `SubGraphMetadata.description`, when displaying
metadata and generating code.
### Supported Input / Output types
TensorFlow Lite metadata for input and output are not designed with specific
model types in mind but rather input and output types. It does not matter what
the model functionally does, as long as the input and output types consists of
the following or a combination of the following, it is supported by TensorFlow
Lite metadata:
* Feature - Numbers which are unsigned integers or float32.
* Image - Metadata currently supports RGB and greyscale images.
* Bounding box - Rectangular shape bounding boxes. The schema supports
[a variety of numbering schemes](https://github.com/tensorflow/tflite-support/blob/4cd0551658b6e26030e0ba7fc4d3127152e0d4ae/tensorflow_lite_support/metadata/metadata_schema.fbs#L214).
### Pack the associated files
TensorFlow Lite models may come with different associated files. For example,
natural language models usually have vocab files that map word pieces to word
IDs; classification models may have label files that indicate object categories.
Without the associated files (if there are), a model will not function well.
The associated files can now be bundled with the model through the metadata
Python library. The new TensorFlow Lite model becomes a zip file that contains
both the model and the associated files. It can be unpacked with common zip
tools. This new model format keeps using the same file extension, `.tflite`. It
is compatible with existing TFLite framework and Interpreter. See
[Pack metadata and associated files into the model](#pack-metadata-and-associated-files-into-the-model)
for more details.
The associated file information can be recorded in the metadata. Depending on
the file type and where the file is attached to (i.e. `ModelMetadata`,
`SubGraphMetadata`, and `TensorMetadata`),
[the TensorFlow Lite Android code generator](../../inference_with_metadata/codegen.md)
may apply corresponding pre/post processing automatically to the object. See
[the \<Codegen usage\> section of each associate file type](https://github.com/tensorflow/tflite-support/blob/4cd0551658b6e26030e0ba7fc4d3127152e0d4ae/tensorflow_lite_support/metadata/metadata_schema.fbs#L77-L127)
in the schema for more details.
### Normalization and quantization parameters
Normalization is a common data preprocessing technique in machine learning. The
goal of normalization is to change the values to a common scale, without
distorting differences in the ranges of values.
[Model quantization](https://www.tensorflow.org/lite/performance/model_optimization#model_quantization)
is a technique that allows for reduced precision representations of weights and
optionally, activations for both storage and computation.
In terms of preprocessing and post-processing, normalization and quantization
are two independent steps. Here are the details.
| | Normalization | Quantization |
| :---------------------: | ----------------------- | ------------------------ |
| \ | **Float model**: \ | **Float model**: \ |
: An example of the : - mean\: 127.5 \ : - zeroPoint\: 0 \ :
: parameter values of the : - std\: 127.5 \ : - scale\: 1.0 \ :
: input image in : **Quant model**\: \ : **Quant model**\: \ :
: MobileNet for float and : - mean\: 127.5 \ : - zeroPoint\: 128.0 \ :
: quant models, : - std\: 127.5 : - scale\:0.0078125f \ :
: respectively. : : :
| \ | \ | **Float models** does |
: \ : \ : not need quantization. \ :
: \ : **Inputs**\: If input : **Quantized model** may :
: \ : data is normalized in : or may not need :
: When to invoke? : training, the input : quantization in pre/post :
: : data of inference needs : processing. It depends :
: : to be normalized : on the datatype of :
: : accordingly. \ : input/output tensors. \ :
: : **Outputs**\: output : - float tensors\: no :
: : data will not be : quantization in pre/post :
: : normalized in general. : processing needed. Quant :
: : : op and dequant op are :
: : : baked into the model :
: : : graph. \ :
: : : - int8/uint8 tensors\: :
: : : need quantization in :
: : : pre/post processing. :
| \ | \ | **Quantize for inputs**: |
: \ : \ : \ :
: Formula : normalized_input = : q = f / scale + :
: : (input - mean) / std : zeroPoint \ :
: : : **Dequantize for :
: : : outputs**\: \ :
: : : f = (q - zeroPoint) * :
: : : scale :
| \ | Filled by model creator | Filled automatically by |
: Where are the : and stored in model : TFLite converter, and :
: parameters : metadata, as : stored in tflite model :
: : `NormalizationOptions` : file. :
| How to get the | Through the | Through the TFLite |
: parameters? : `MetadataExtractor` API : `Tensor` API [1] or :
: : [2] : through the :
: : : `MetadataExtractor` API :
: : : [2] :
| Do float and quant | Yes, float and quant | No, the float model does |
: models share the same : models have the same : not need quantization. :
: value? : Normalization : :
: : parameters : :
| Does TFLite Code | \ | \ |
: generator or Android : Yes : Yes :
: Studio ML binding : : :
: automatically generate : : :
: it in data processing? : : :
[1] The
[TensorFlow Lite Java API](https://github.com/tensorflow/tensorflow/blob/09ec15539eece57b257ce9074918282d88523d56/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Tensor.java#L73)
and the
[TensorFlow Lite C++ API](https://github.com/tensorflow/tensorflow/blob/09ec15539eece57b257ce9074918282d88523d56/tensorflow/lite/c/common.h#L391).
\
[2] The [metadata extractor library](#read-the-metadata-from-models)
When processing image data for uint8 models, normalization and quantization are
sometimes skipped. It is fine to do so when the pixel values are in the range of
[0, 255]. But in general, you should always process the data according to the
normalization and quantization parameters when applicable.
[TensorFlow Lite Task Library](https://www.tensorflow.org/lite/inference_with_metadata/overview)
can handle normalization for you if you set up `NormalizationOptions` in
metadata. Quantization and dequantization processing is always encapsulated.
### Examples
Note: The export directory specified has to exist before you run the script; it
does not get created as part of the process.
You can find examples on how the metadata should be populated for different
types of models here:
#### Image classification
Download the script
[here](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/metadata/metadata_writer_for_image_classifier.py)
, which populates metadata to
[mobilenet_v1_0.75_160_quantized.tflite](https://tfhub.dev/tensorflow/lite-model/mobilenet_v1_0.75_160_quantized/1/default/1).
Run the script like this:
```sh
python ./metadata_writer_for_image_classifier.py \
--model_file=./model_without_metadata/mobilenet_v1_0.75_160_quantized.tflite \
--label_file=./model_without_metadata/labels.txt \
--export_directory=model_with_metadata
```
To populate metadata for other image classification models, add the model specs
like
[this](https://github.com/tensorflow/examples/blob/master/lite/examples/image_classification/metadata/metadata_writer_for_image_classifier.py#L63-L74)
into the script. The rest of this guide will highlight some of the key sections
in the image classification example to illustrate the key elements.
### Deep dive into the image classification example
#### Model information
Metadata starts by creating a new model info:
```python
from tflite_support import flatbuffers
from tflite_support import metadata as _metadata
from tflite_support import metadata_schema_py_generated as _metadata_fb
""" ... """
"""Creates the metadata for an image classifier."""
# Creates model info.
model_meta = _metadata_fb.ModelMetadataT()
model_meta.name = "MobileNetV1 image classifier"
model_meta.description = ("Identify the most prominent object in the "
"image from a set of 1,001 categories such as "
"trees, animals, food, vehicles, person etc.")
model_meta.version = "v1"
model_meta.author = "TensorFlow"
model_meta.license = ("Apache License. Version 2.0 "
"http://www.apache.org/licenses/LICENSE-2.0.")
```
#### Input / output information
This section shows you how to describe your model's input and output signature.
This metadata may be used by automatic code generators to create pre- and post-
processing code. To create input or output information about a tensor:
```python
# Creates input info.
input_meta = _metadata_fb.TensorMetadataT()
# Creates output info.
output_meta = _metadata_fb.TensorMetadataT()
```
#### Image input
Image is a common input type for machine learning. TensorFlow Lite metadata
supports information such as colorspace and pre-processing information such as
normalization. The dimension of the image does not require manual specification
since it is already provided by the shape of the input tensor and can be
automatically inferred.
```python
input_meta.name = "image"
input_meta.description = (
"Input image to be classified. The expected image is {0} x {1}, with "
"three channels (red, blue, and green) per pixel. Each value in the "
"tensor is a single byte between 0 and 255.".format(160, 160))
input_meta.content = _metadata_fb.ContentT()
input_meta.content.contentProperties = _metadata_fb.ImagePropertiesT()
input_meta.content.contentProperties.colorSpace = (
_metadata_fb.ColorSpaceType.RGB)
input_meta.content.contentPropertiesType = (
_metadata_fb.ContentProperties.ImageProperties)
input_normalization = _metadata_fb.ProcessUnitT()
input_normalization.optionsType = (
_metadata_fb.ProcessUnitOptions.NormalizationOptions)
input_normalization.options = _metadata_fb.NormalizationOptionsT()
input_normalization.options.mean = [127.5]
input_normalization.options.std = [127.5]
input_meta.processUnits = [input_normalization]
input_stats = _metadata_fb.StatsT()
input_stats.max = [255]
input_stats.min = [0]
input_meta.stats = input_stats
```
#### Label output
Label can be mapped to an output tensor via an associated file using
`TENSOR_AXIS_LABELS`.
```python
# Creates output info.
output_meta = _metadata_fb.TensorMetadataT()
output_meta.name = "probability"
output_meta.description = "Probabilities of the 1001 labels respectively."
output_meta.content = _metadata_fb.ContentT()
output_meta.content.content_properties = _metadata_fb.FeaturePropertiesT()
output_meta.content.contentPropertiesType = (
_metadata_fb.ContentProperties.FeatureProperties)
output_stats = _metadata_fb.StatsT()
output_stats.max = [1.0]
output_stats.min = [0.0]
output_meta.stats = output_stats
label_file = _metadata_fb.AssociatedFileT()
label_file.name = os.path.basename("your_path_to_label_file")
label_file.description = "Labels for objects that the model can recognize."
label_file.type = _metadata_fb.AssociatedFileType.TENSOR_AXIS_LABELS
output_meta.associatedFiles = [label_file]
```
#### Create the metadata Flatbuffers
The following code combines the model information with the input and output
information:
```python
# Creates subgraph info.
subgraph = _metadata_fb.SubGraphMetadataT()
subgraph.inputTensorMetadata = [input_meta]
subgraph.outputTensorMetadata = [output_meta]
model_meta.subgraphMetadata = [subgraph]
b = flatbuffers.Builder(0)
b.Finish(
model_meta.Pack(b),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
metadata_buf = b.Output()
```
#### Pack metadata and associated files into the model
Once the metadata Flatbuffers is created, the metadata and the label file are
written into the TFLite file via the `populate` method:
```python
populator = _metadata.MetadataPopulator.with_model_file(model_file)
populator.load_metadata_buffer(metadata_buf)
populator.load_associated_files(["your_path_to_label_file"])
populator.populate()
```
You can pack as many associated files as you want into the model through
`load_associated_files`. However, it is required to pack at least those files
documented in the metadata. In this example, packing the label file is
mandatory.
## Visualize the metadata
You can use [Netron](https://github.com/lutzroeder/netron) to visualize your
metadata, or you can read the metadata from a TensorFlow Lite model into a json
format using the `MetadataDisplayer`:
```python
displayer = _metadata.MetadataDisplayer.with_model_file(export_model_path)
export_json_file = os.path.join(FLAGS.export_directory,
os.path.splitext(model_basename)[0] + ".json")
json_file = displayer.get_metadata_json()
# Optional: write out the metadata as a json file
with open(export_json_file, "w") as f:
f.write(json_file)
```
Android Studio also supports displaying metadata through the
[Android Studio ML Binding feature](https://developer.android.com/studio/preview/features#tensor-flow-lite-models).
## Metadata versioning
The
[metadata schema](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/metadata_schema.fbs)
is versioned both by the Semantic versioning number, which tracks the changes of
the schema file, and by the Flatbuffers file identification, which indicates the
true version compatibility.
### The Semantic versioning number
The metadata schema is versioned by the
[Semantic versioning number](https://github.com/tensorflow/tflite-support/blob/4cd0551658b6e26030e0ba7fc4d3127152e0d4ae/tensorflow_lite_support/metadata/metadata_schema.fbs#L53),
such as MAJOR.MINOR.PATCH. It tracks schema changes according to the rules
[here](https://github.com/tensorflow/tflite-support/blob/4cd0551658b6e26030e0ba7fc4d3127152e0d4ae/tensorflow_lite_support/metadata/metadata_schema.fbs#L32-L44).
See the
[history of fields](https://github.com/tensorflow/tflite-support/blob/4cd0551658b6e26030e0ba7fc4d3127152e0d4ae/tensorflow_lite_support/metadata/metadata_schema.fbs#L63)
added after version `1.0.0`.
### The Flatbuffers file identification
Semantic versioning guarantees the compatibility if following the rules, but it
does not imply the true incompatibility. When bumping up the MAJOR number, it
does not necessarily mean the backward compatibility is broken. Therefore, we
use the
[Flatbuffers file identification](https://google.github.io/flatbuffers/md__schemas.html),
[file_identifier](https://github.com/tensorflow/tflite-support/blob/4cd0551658b6e26030e0ba7fc4d3127152e0d4ae/tensorflow_lite_support/metadata/metadata_schema.fbs#L61),
to denote the true compatibility of the metadata schema. The file identifier is
exactly 4 characters long. It is fixed to a certain metadata schema and not
subject to change by users. If the backward compatibility of the metadata schema
has to be broken for some reason, the file_identifier will bump up, for example,
from “M001” to “M002”. File_identifier is expected to be changed much less
frequently than the metadata_version.
### The minimum necessary metadata parser version
The
[minimum necessary metadata parser version](https://github.com/tensorflow/tflite-support/blob/4cd0551658b6e26030e0ba7fc4d3127152e0d4ae/tensorflow_lite_support/metadata/metadata_schema.fbs#L681)
is the minimum version of metadata parser (the Flatbuffers generated code) that
can read the metadata Flatbuffers in full. The version is effectively the
largest version number among the versions of all the fields populated and the
smallest compatible version indicated by the file identifier. The minimum
necessary metadata parser version is automatically populated by the
`MetadataPopulator` when the metadata is populated into a TFLite model. See the
[metadata extractor](#read-the-metadata-from-models) for more information on how
the minimum necessary metadata parser version is used.
## Read the metadata from models
The Metadata Extractor library is convenient tool to read the metadata and
associated files from a models across different platforms (see the
[Java version](https://github.com/tensorflow/tflite-support/tree/master/tensorflow_lite_support/metadata/java)
and the
[C++ version](https://github.com/tensorflow/tflite-support/tree/master/tensorflow_lite_support/metadata/cc)).
You can build your own metadata extractor tool in other languages using the
Flatbuffers library.
### Read the metadata in Java
To use the Metadata Extractor library in your Android app, we recommend using
the
[TensorFlow Lite Metadata AAR hosted at MavenCentral](https://search.maven.org/artifact/org.tensorflow/tensorflow-lite-metadata).
It contains the `MetadataExtractor` class, as well as the FlatBuffers Java
bindings for the
[metadata schema](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/metadata_schema.fbs)
and the
[model schema](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/schema/schema.fbs).
You can specify this in your `build.gradle` dependencies as follows:
```build
dependencies {
implementation 'org.tensorflow:tensorflow-lite-metadata:0.1.0'
}
```
To use nightly snapshots, make sure that you have added
[Sonatype snapshot repository](https://www.tensorflow.org/lite/android/lite_build#use_nightly_snapshots).
You can initialize a `MetadataExtractor` object with a `ByteBuffer` that points
to the model:
```java
public MetadataExtractor(ByteBuffer buffer);
```
The `ByteBuffer` must remain unchanged for the entire lifetime of the
`MetadataExtractor` object. The initialization may fail if the Flatbuffers file
identifier of the model metadata does not match that of the metadata parser. See
[metadata versioning](#metadata-versioning) for more information.
With matching file identifiers, the metadata extractor will successfully read
metadata generated from all past and future schema due to the Flatbuffers'
forwards and backward compatibility mechanism. However, fields from future
schemas cannot be extracted by older metadata extractors. The
[minimum necessary parser version](#the-minimum-necessary-metadata-parser-version)
of the metadata indicates the minimum version of metadata parser that can read
the metadata Flatbuffers in full. You can use the following method to verify if
the minimum necessary parser version condition is met:
```java
public final boolean isMinimumParserVersionSatisfied();
```
Passing in a model without metadata is allowed. However, invoking methods that
read from the metadata will cause runtime errors. You can check if a model has
metadata by invoking the `hasMetadata` method:
```java
public boolean hasMetadata();
```
`MetadataExtractor` provides convenient functions for you to get the
input/output tensors' metadata. For example,
```java
public int getInputTensorCount();
public TensorMetadata getInputTensorMetadata(int inputIndex);
public QuantizationParams getInputTensorQuantizationParams(int inputIndex);
public int[] getInputTensorShape(int inputIndex);
public int getoutputTensorCount();
public TensorMetadata getoutputTensorMetadata(int inputIndex);
public QuantizationParams getoutputTensorQuantizationParams(int inputIndex);
public int[] getoutputTensorShape(int inputIndex);
```
Though the
[TensorFlow Lite model schema](https://github.com/tensorflow/tensorflow/blob/aa7ff6aa28977826e7acae379e82da22482b2bf2/tensorflow/lite/schema/schema.fbs#L1075)
supports multiple subgraphs, the TFLite Interpreter currently only supports a
single subgraph. Therefore, `MetadataExtractor` omits subgraph index as an input
argument in its methods.
## Read the associated files from models
The TensorFlow Lite model with metadata and associated files is essentially a
zip file that can be unpacked with common zip tools to get the associated files.
For example, you can unzip
[mobilenet_v1_0.75_160_quantized](https://tfhub.dev/tensorflow/lite-model/mobilenet_v1_0.75_160_quantized/1/metadata/1)
and extract the label file in the model as follows:
```sh
$ unzip mobilenet_v1_0.75_160_quantized_1_metadata_1.tflite
Archive: mobilenet_v1_0.75_160_quantized_1_metadata_1.tflite
extracting: labels.txt
```
You can also read associated files through the Metadata Extractor library.
In Java, pass the file name into the `MetadataExtractor.getAssociatedFile`
method:
```java
public InputStream getAssociatedFile(String fileName);
```
Similarly, in C++, this can be done with the method,
`ModelMetadataExtractor::GetAssociatedFile`:
```c++
tflite::support::StatusOr<absl::string_view> GetAssociatedFile(
const std::string& filename) const;
```
@@ -0,0 +1,847 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "Mq-riZs-TJGt"
},
"source": [
"##### Copyright 2024 The AI Edge Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "LEvnopDoTC4M"
},
"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": "QSRG6qmtTRSk"
},
"source": [
"# TensorFlow Lite Metadata Writer API\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JlzjEt4Txr0x"
},
"source": [
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://www.tensorflow.org/lite/models/convert/metadata_writer_tutorial\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/convert/metadata_writer_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/convert/metadata_writer_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/models/convert/metadata_writer_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
" </td>\n",
"\n",
"</table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "b0gwEhfRYat6"
},
"source": [
"[TensorFlow Lite Model Metadata](https://www.tensorflow.org/lite/models/convert/metadata) is a standard model description format. It contains rich semantics for general model information, inputs/outputs, and associated files, which makes the model self-descriptive and exchangeable.\n",
"\n",
"Model Metadata is currently used in the following two primary use cases:\n",
"1. **Enable easy model inference using TensorFlow Lite [Task Library](https://www.tensorflow.org/lite/inference_with_metadata/task_library/overview) and [codegen tools](https://www.tensorflow.org/lite/inference_with_metadata/codegen)**. Model Metadata contains the mandatory information required during inference, such as label files in image classification, sampling rate of the audio input in audio classification, and tokenizer type to process input string in Natural Language models.\n",
"\n",
"2. **Enable model creators to include documentation**, such as the description of model inputs/outputs or how to use the model. Model users can view this documentation via visualization tools such as [Netron](https://netron.app/).\n",
"\n",
"TensorFlow Lite Metadata Writer API provides an easy-to-use API to create Model Metadata for popular ML tasks supported by the TFLite Task Library. This notebook shows examples of how the metadata should be populated for the following tasks below:\n",
"\n",
"* [Image classifiers](#image_classifiers)\n",
"* [Object detectors](#object_detectors)\n",
"* [Image segmenters](#image_segmenters)\n",
"* [Natural language classifiers](#nl_classifiers)\n",
"* [Audio classifiers](#audio_classifiers)\n",
"\n",
"Metadata writers for BERT natural language classifiers and BERT question answerers are coming soon.\n",
"\n",
"If you want to add metadata for use cases that are not supported, please use the [Flatbuffers Python API](https://www.tensorflow.org/lite/models/convert/metadata#adding_metadata). See the tutorials [here](https://www.tensorflow.org/lite/models/convert/metadata#adding_metadata).\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GVRIGdA4T6tO"
},
"source": [
"## Prerequisites"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bVTD2KSyotBK"
},
"source": [
"Install the TensorFlow Lite Support Pypi package."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "m-8xSrSvUg-6"
},
"outputs": [],
"source": [
"!pip install tflite-support-nightly"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hyYS87Odpxef"
},
"source": [
"## Create Model Metadata for Task Library and Codegen"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uLxv541TqTim"
},
"source": [
"<a name=image_classifiers></a>\n",
"### Image classifiers"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "s41TjCGlsyEF"
},
"source": [
"See the [image classifier model compatibility requirements](https://www.tensorflow.org/lite/inference_with_metadata/task_library/image_classifier#model_compatibility_requirements) for more details about the supported model format."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_KsPKmg8T9-8"
},
"source": [
"Step 1: Import the required packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "hhgNqEtWrwB3"
},
"outputs": [],
"source": [
"from tflite_support.metadata_writers import image_classifier\n",
"from tflite_support.metadata_writers import writer_utils"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "o9WBgiFdsiIQ"
},
"source": [
"Step 2: Download the example image classifier, [mobilenet_v2_1.0_224.tflite](https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/mobilenet_v2_1.0_224.tflite), and the [label file](https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/labels.txt)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6WgSBbNet-Tt"
},
"outputs": [],
"source": [
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/mobilenet_v2_1.0_224.tflite -o mobilenet_v2_1.0_224.tflite\n",
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/labels.txt -o mobilenet_labels.txt"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ALtlz7woweHe"
},
"source": [
"Step 3: Create metadata writer and populate."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "_SMEBBt2r-W6"
},
"outputs": [],
"source": [
"ImageClassifierWriter = image_classifier.MetadataWriter\n",
"_MODEL_PATH = \"mobilenet_v2_1.0_224.tflite\"\n",
"# Task Library expects label files that are in the same format as the one below.\n",
"_LABEL_FILE = \"mobilenet_labels.txt\"\n",
"_SAVE_TO_PATH = \"mobilenet_v2_1.0_224_metadata.tflite\"\n",
"# Normalization parameters are required when reprocessing the image. It is\n",
"# optional if the image pixel values are in the range of [0, 255] and the input\n",
"# tensor is quantized to uint8. See the introduction for normalization and\n",
"# quantization parameters below for more details.\n",
"# https://www.tensorflow.org/lite/models/convert/metadata#normalization_and_quantization_parameters)\n",
"_INPUT_NORM_MEAN = 127.5\n",
"_INPUT_NORM_STD = 127.5\n",
"\n",
"# Create the metadata writer.\n",
"writer = ImageClassifierWriter.create_for_inference(\n",
" writer_utils.load_file(_MODEL_PATH), [_INPUT_NORM_MEAN], [_INPUT_NORM_STD],\n",
" [_LABEL_FILE])\n",
"\n",
"# Verify the metadata generated by the metadata writer.\n",
"print(writer.get_metadata_json())\n",
"\n",
"# Populate the metadata into the model.\n",
"writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GhhTDkr-uf0n"
},
"source": [
"<a name=object_detectors></a>\n",
"### Object detectors"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EL9GssnTuf0n"
},
"source": [
"See the [object detector model compatibility requirements](https://www.tensorflow.org/lite/inference_with_metadata/task_library/object_detector#model_compatibility_requirements) for more details about the supported model format."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "r-HUTEtHuf0n"
},
"source": [
"Step 1: Import the required packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2_NIROeouf0o"
},
"outputs": [],
"source": [
"from tflite_support.metadata_writers import object_detector\n",
"from tflite_support.metadata_writers import writer_utils"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "UM6jijiUuf0o"
},
"source": [
"Step 2: Download the example object detector, [ssd_mobilenet_v1.tflite](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/object_detector/ssd_mobilenet_v1.tflite), and the [label file](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/object_detector/labelmap.txt)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4i_BBfGzuf0o"
},
"outputs": [],
"source": [
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/object_detector/ssd_mobilenet_v1.tflite -o ssd_mobilenet_v1.tflite\n",
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/object_detector/labelmap.txt -o ssd_mobilenet_labels.txt"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "DG9T3eSDwsnd"
},
"source": [
"Step 3: Create metadata writer and populate."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "vMGGeJfCuf0p"
},
"outputs": [],
"source": [
"ObjectDetectorWriter = object_detector.MetadataWriter\n",
"_MODEL_PATH = \"ssd_mobilenet_v1.tflite\"\n",
"# Task Library expects label files that are in the same format as the one below.\n",
"_LABEL_FILE = \"ssd_mobilenet_labels.txt\"\n",
"_SAVE_TO_PATH = \"ssd_mobilenet_v1_metadata.tflite\"\n",
"# Normalization parameters are required when reprocessing the image. It is\n",
"# optional if the image pixel values are in the range of [0, 255] and the input\n",
"# tensor is quantized to uint8. See the introduction for normalization and\n",
"# quantization parameters below for more details.\n",
"# https://www.tensorflow.org/lite/models/convert/metadata#normalization_and_quantization_parameters)\n",
"_INPUT_NORM_MEAN = 127.5\n",
"_INPUT_NORM_STD = 127.5\n",
"\n",
"# Create the metadata writer.\n",
"writer = ObjectDetectorWriter.create_for_inference(\n",
" writer_utils.load_file(_MODEL_PATH), [_INPUT_NORM_MEAN], [_INPUT_NORM_STD],\n",
" [_LABEL_FILE])\n",
"\n",
"# Verify the metadata generated by the metadata writer.\n",
"print(writer.get_metadata_json())\n",
"\n",
"# Populate the metadata into the model.\n",
"writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "QT0Oa0SU6uGS"
},
"source": [
"<a name=image_segmenters></a>\n",
"### Image segmenters"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "XaFQmg-S6uGW"
},
"source": [
"See the [image segmenter model compatibility requirements](https://www.tensorflow.org/lite/inference_with_metadata/task_library/image_segmenter#model_compatibility_requirements) for more details about the supported model format."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "DiktANhj6uGX"
},
"source": [
"Step 1: Import the required packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "H6Lrw3op6uGX"
},
"outputs": [],
"source": [
"from tflite_support.metadata_writers import image_segmenter\n",
"from tflite_support.metadata_writers import writer_utils"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9EFs8Oyi6uGX"
},
"source": [
"Step 2: Download the example image segmenter, [deeplabv3.tflite](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/image_segmenter/deeplabv3.tflite), and the [label file](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/image_segmenter/labelmap.txt)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "feQDH0bN6uGY"
},
"outputs": [],
"source": [
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_segmenter/deeplabv3.tflite -o deeplabv3.tflite\n",
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_segmenter/labelmap.txt -o deeplabv3_labels.txt"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8LhiAbJM6uGY"
},
"source": [
"Step 3: Create metadata writer and populate."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yot8xLI46uGY"
},
"outputs": [],
"source": [
"ImageSegmenterWriter = image_segmenter.MetadataWriter\n",
"_MODEL_PATH = \"deeplabv3.tflite\"\n",
"# Task Library expects label files that are in the same format as the one below.\n",
"_LABEL_FILE = \"deeplabv3_labels.txt\"\n",
"_SAVE_TO_PATH = \"deeplabv3_metadata.tflite\"\n",
"# Normalization parameters are required when reprocessing the image. It is\n",
"# optional if the image pixel values are in the range of [0, 255] and the input\n",
"# tensor is quantized to uint8. See the introduction for normalization and\n",
"# quantization parameters below for more details.\n",
"# https://www.tensorflow.org/lite/models/convert/metadata#normalization_and_quantization_parameters)\n",
"_INPUT_NORM_MEAN = 127.5\n",
"_INPUT_NORM_STD = 127.5\n",
"\n",
"# Create the metadata writer.\n",
"writer = ImageSegmenterWriter.create_for_inference(\n",
" writer_utils.load_file(_MODEL_PATH), [_INPUT_NORM_MEAN], [_INPUT_NORM_STD],\n",
" [_LABEL_FILE])\n",
"\n",
"# Verify the metadata generated by metadata writer.\n",
"print(writer.get_metadata_json())\n",
"\n",
"# Populate the metadata into the model.\n",
"writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "NnvM80e7AG-h"
},
"source": [
"<a name=nl_classifiers></a>\n",
"###Natural language classifiers"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dfOPhFwOAG-k"
},
"source": [
"See the [natural language classifier model compatibility requirements](https://www.tensorflow.org/lite/inference_with_metadata/task_library/nl_classifier#model_compatibility_requirements) for more details about the supported model format."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "WMJ7tvuwAG-k"
},
"source": [
"Step 1: Import the required packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "_FGVyb2iAG-k"
},
"outputs": [],
"source": [
"from tflite_support.metadata_writers import nl_classifier\n",
"from tflite_support.metadata_writers import metadata_info\n",
"from tflite_support.metadata_writers import writer_utils"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "iIg7rATpAG-l"
},
"source": [
"Step 2: Download the example natural language classifier, [movie_review.tflite](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/nl_classifier/movie_review.tflite), the [label file](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/nl_classifier/labels.txt), and the [vocab file](https://storage.googleapis.com/download.tensorflow.org/models/tflite_support/nl_classifier/vocab.txt)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "TzuQcti2AG-l"
},
"outputs": [],
"source": [
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/nl_classifier/movie_review.tflite -o movie_review.tflite\n",
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/nl_classifier/labels.txt -o movie_review_labels.txt\n",
"!curl -L https://storage.googleapis.com/download.tensorflow.org/models/tflite_support/nl_classifier/vocab.txt -o movie_review_vocab.txt"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BWxUtHdeAG-m"
},
"source": [
"Step 3: Create metadata writer and populate."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "NGPWzRuHAG-m"
},
"outputs": [],
"source": [
"NLClassifierWriter = nl_classifier.MetadataWriter\n",
"_MODEL_PATH = \"movie_review.tflite\"\n",
"# Task Library expects label files and vocab files that are in the same formats\n",
"# as the ones below.\n",
"_LABEL_FILE = \"movie_review_labels.txt\"\n",
"_VOCAB_FILE = \"movie_review_vocab.txt\"\n",
"# NLClassifier supports tokenize input string using the regex tokenizer. See\n",
"# more details about how to set up RegexTokenizer below:\n",
"# https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/metadata_writers/metadata_info.py#L130\n",
"_DELIM_REGEX_PATTERN = r\"[^\\w\\']+\"\n",
"_SAVE_TO_PATH = \"moview_review_metadata.tflite\"\n",
"\n",
"# Create the metadata writer.\n",
"writer = nl_classifier.MetadataWriter.create_for_inference(\n",
" writer_utils.load_file(_MODEL_PATH),\n",
" metadata_info.RegexTokenizerMd(_DELIM_REGEX_PATTERN, _VOCAB_FILE),\n",
" [_LABEL_FILE])\n",
"\n",
"# Verify the metadata generated by metadata writer.\n",
"print(writer.get_metadata_json())\n",
"\n",
"# Populate the metadata into the model.\n",
"writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qv0WDnzW711f"
},
"source": [
"<a name=audio_classifiers></a>\n",
"### Audio classifiers"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xqP7X8jww8pL"
},
"source": [
"See the [audio classifier model compatibility requirements](https://www.tensorflow.org/lite/inference_with_metadata/task_library/audio_classifier#model_compatibility_requirements) for more details about the supported model format."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7RToKepxw8pL"
},
"source": [
"Step 1: Import the required packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "JjddvTXKw8pL"
},
"outputs": [],
"source": [
"from tflite_support.metadata_writers import audio_classifier\n",
"from tflite_support.metadata_writers import metadata_info\n",
"from tflite_support.metadata_writers import writer_utils"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ar418rH6w8pL"
},
"source": [
"Step 2: Download the example audio classifier, [yamnet.tflite](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/audio_classifier/yamnet_wavin_quantized_mel_relu6.tflite), and the [label file](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/audio_classifier/yamnet_521_labels.txt)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5eQY6znmw8pM"
},
"outputs": [],
"source": [
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/audio_classifier/yamnet_wavin_quantized_mel_relu6.tflite -o yamnet.tflite\n",
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/audio_classifier/yamnet_521_labels.txt -o yamnet_labels.txt\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "1TYP5w0Ew8pM"
},
"source": [
"Step 3: Create metadata writer and populate."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "MDlSczBQw8pM"
},
"outputs": [],
"source": [
"AudioClassifierWriter = audio_classifier.MetadataWriter\n",
"_MODEL_PATH = \"yamnet.tflite\"\n",
"# Task Library expects label files that are in the same format as the one below.\n",
"_LABEL_FILE = \"yamnet_labels.txt\"\n",
"# Expected sampling rate of the input audio buffer.\n",
"_SAMPLE_RATE = 16000\n",
"# Expected number of channels of the input audio buffer. Note that Task library only\n",
"# supports single channel so far.\n",
"_CHANNELS = 1\n",
"_SAVE_TO_PATH = \"yamnet_metadata.tflite\"\n",
"\n",
"# Create the metadata writer.\n",
"writer = AudioClassifierWriter.create_for_inference(\n",
" writer_utils.load_file(_MODEL_PATH), _SAMPLE_RATE, _CHANNELS, [_LABEL_FILE])\n",
"\n",
"# Verify the metadata generated by the metadata writer.\n",
"print(writer.get_metadata_json())\n",
"\n",
"# Populate the metadata into the model.\n",
"writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "YoRLs84yNAJR"
},
"source": [
"## Create Model Metadata with semantic information"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cxXsOBknOGJ2"
},
"source": [
"You can fill in more descriptive information about the model and each tensor through the Metadata Writer API to help improve model understanding. It can be done through the 'create_from_metadata_info' method in each metadata writer. In general, you can fill in data through the parameters of 'create_from_metadata_info', i.e. `general_md`, `input_md`, and `output_md`. See the example below to create a rich Model Metadata for image classifiers."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Q-LW6nrcQ9lv"
},
"source": [
"Step 1: Import the required packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "KsL_egYcRGw3"
},
"outputs": [],
"source": [
"from tflite_support.metadata_writers import image_classifier\n",
"from tflite_support.metadata_writers import metadata_info\n",
"from tflite_support.metadata_writers import writer_utils\n",
"from tflite_support import metadata_schema_py_generated as _metadata_fb"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0UWck_8uRboF"
},
"source": [
"Step 2: Download the example image classifier, [mobilenet_v2_1.0_224.tflite](https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/mobilenet_v2_1.0_224.tflite), and the [label file](https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/labels.txt)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "TqJ-jh-PRVdk"
},
"outputs": [],
"source": [
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/mobilenet_v2_1.0_224.tflite -o mobilenet_v2_1.0_224.tflite\n",
"!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/labels.txt -o mobilenet_labels.txt"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "r4I5wJMQRxzb"
},
"source": [
"Step 3: Create model and tensor information."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "urd7HDuaR_HC"
},
"outputs": [],
"source": [
"model_buffer = writer_utils.load_file(\"mobilenet_v2_1.0_224.tflite\")\n",
"\n",
"# Create general model information.\n",
"general_md = metadata_info.GeneralMd(\n",
" name=\"ImageClassifier\",\n",
" version=\"v1\",\n",
" description=(\"Identify the most prominent object in the image from a \"\n",
" \"known set of categories.\"),\n",
" author=\"TensorFlow Lite\",\n",
" licenses=\"Apache License. Version 2.0\")\n",
"\n",
"# Create input tensor information.\n",
"input_md = metadata_info.InputImageTensorMd(\n",
" name=\"input image\",\n",
" description=(\"Input image to be classified. The expected image is \"\n",
" \"128 x 128, with three channels (red, blue, and green) per \"\n",
" \"pixel. Each element in the tensor is a value between min and \"\n",
" \"max, where (per-channel) min is [0] and max is [255].\"),\n",
" norm_mean=[127.5],\n",
" norm_std=[127.5],\n",
" color_space_type=_metadata_fb.ColorSpaceType.RGB,\n",
" tensor_type=writer_utils.get_input_tensor_types(model_buffer)[0])\n",
"\n",
"# Create output tensor information.\n",
"output_md = metadata_info.ClassificationTensorMd(\n",
" name=\"probability\",\n",
" description=\"Probabilities of the 1001 labels respectively.\",\n",
" label_files=[\n",
" metadata_info.LabelFileMd(file_path=\"mobilenet_labels.txt\",\n",
" locale=\"en\")\n",
" ],\n",
" tensor_type=writer_utils.get_output_tensor_types(model_buffer)[0])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "N5aL5Uxkf4aO"
},
"source": [
"Step 4: Create metadata writer and populate."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "_iWIwdqEf_mr"
},
"outputs": [],
"source": [
"ImageClassifierWriter = image_classifier.MetadataWriter\n",
"# Create the metadata writer.\n",
"writer = ImageClassifierWriter.create_from_metadata_info(\n",
" model_buffer, general_md, input_md, output_md)\n",
"\n",
"# Verify the metadata generated by the metadata writer.\n",
"print(writer.get_metadata_json())\n",
"\n",
"# Populate the metadata into the model.\n",
"writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "z78vuu6np5sb"
},
"source": [
"## Read the metadata populated to your model."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "DnWt-4oOselo"
},
"source": [
"You can display the metadata and associated files in a TFLite model through the following code:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5D13YPUsp5VT"
},
"outputs": [],
"source": [
"from tflite_support import metadata\n",
"\n",
"displayer = metadata.MetadataDisplayer.with_model_file(\"mobilenet_v2_1.0_224_metadata.tflite\")\n",
"print(\"Metadata populated:\")\n",
"print(displayer.get_metadata_json())\n",
"\n",
"print(\"Associated file(s) populated:\")\n",
"for file_name in displayer.get_packed_associated_file_list():\n",
" print(\"file name: \", file_name)\n",
" print(\"file content:\")\n",
" print(displayer.get_associated_file_buffer(file_name))"
]
}
],
"metadata": {
"colab": {
"collapsed_sections": [
"Mq-riZs-TJGt"
],
"name": "Metadata Writer tutorial",
"private_outputs": true,
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
@@ -0,0 +1,348 @@
# TensorFlow operation fusion
## Overview
This page describes the design and steps needed to convert composite operations
in TensorFlow to fused operations in TensorFlow Lite. This infrastructure is
general purpose and supports conversion of any composite operation in TensorFlow
to a corresponding fused operation in TensorFlow Lite.
An example use of this infrastructure is TensorFlow RNN operation fusion to
TensorFlow Lite, as detailed
[here](https://www.tensorflow.org/lite/models/convert/rnn).
### What are fused operations
![drawing](../../images/convert/op_fusion_banner.jpg)
TensorFlow operations can either be primitive ops e.g.
[tf.add](https://www.tensorflow.org/api_docs/python/tf/math/add) or they can be
composed from other primitive operations e.g.
[tf.einsum](https://www.tensorflow.org/api_docs/python/tf/einsum). A primitive
operation shows up as a single node in the TensorFlow graph while.a composite
operation is a collection of nodes in the TensorFlow graph. Executing a
composite operation is equivalent to executing each of its constituent primitive
operations.
A fused operation corresponds to a single operation that subsumes all the
computation performed by each primitive operation within the corresponding
composite operation.
### Benefits of fused operations
Fused operations exist to maximize the performance of their underlying kernel
implementations, by optimizing the overall computation and reducing memory
footprint. This is very valuable, especially for low-latency inference workloads
and resource constrained mobile platforms.
Fused operations also provide a higher level interface to define complex
transformations like quantization, which would otherwise be infeasible or very
hard to do at a more granular level.
TensorFlow Lite has many instances of fused operations for the reasons
articulated above. These fused operations typically correspond to composite
operations in the source TensorFlow program. Examples of composite operations in
TensorFlow that are implemented as a single fused operation in TensorFlow Lite
include various RNN operations like Unidirectional and Bidirectional sequence
LSTM, convolution (conv2d, bias add, relu), fully connected (matmul, bias add,
relu) and more. In TensorFlow Lite, LSTM quantization is currently only
implemented in the fused LSTM operations.
### Challenges with fused operations
Converting composite operations from TensorFlow to fused operations in
TensorFlow Lite is a hard problem. This is because:
1. Composite operations are represented in the TensorFlow graph as a set of
primitive operations without a well defined boundary. It can be very
challenging to identify (e.g. via pattern matching) the sub-graph
corresponding to such a composite operation.
1. There may be more than one TensorFlow implementation targeting a fused
TensorFlow Lite operation. For example, there are many LSTM implementations
in TensorFlow (Keras, Babelfish/lingvo etc) and each of these is composed of
different primitive operations but they all could still be converted to the
same fused LSTM operation in TensorFlow Lite.
As such, conversion of fused operations has proven quite challenging.
## Converting from composite op to a TFLite custom operation (recommended)
### Wrap the composite operation in a `tf.function`
In many cases, some part of the model can be mapped to a single operation in
TFLite. This can help with performance when writing an optimized implementation
for specific operations. To be able to create a fused operation in TFLite,
identify the part of the graph that represents a fused operation and wrap it in
a `tf.function` with "experimental_implements" attribute to a `tf.function`,
which has attribute value `tfl_fusable_op` with value `true`. If the custom
operation takes attributes then pass them as part of the same
"experimental_implements".
Example,
```python
def get_implements_signature():
implements_signature = [
# 'name' will be used as a name for the operation.
'name: "my_custom_fused_op"',
# attr "tfl_fusable_op" is required to be set with true value.
'attr {key: "tfl_fusable_op" value { b: true } }',
# Example attribute "example_option" that the op accepts.
'attr {key: "example_option" value { i: %d } }' % 10
]
return ' '.join(implements_signature)
@tf.function(experimental_implements=get_implements_signature())
def my_custom_fused_op(input_1, input_2):
# An empty function that represents pre/post processing example that
# is not represented as part of the Tensorflow graph.
output_1 = tf.constant(0.0, dtype=tf.float32, name='first_output')
output_2 = tf.constant(0.0, dtype=tf.float32, name='second_output')
return output_1, output_2
class TestModel(tf.Module):
def __init__(self):
super(TestModel, self).__init__()
self.conv_1 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
self.conv_2 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
@tf.function(input_signature=[
tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
])
def simple_eval(self, input_a, input_b):
return my_custom_fused_op(self.conv_1(input_a), self.conv_2(input_b))
```
Note that you don't need to set `allow_custom_ops` on the converter as
`tfl_fusable_op` attribute imply this already.
### Implement custom op and register with TFLite Interpreter
Implement your fused operation as a TFLite Custom operation - see
[instructions](https://www.tensorflow.org/lite/guide/ops_custom).
Note that, the name to register the op with should be similar to the name
specified in the `name` attribute in the implements signature.
An example for the op in the example is
```c++
TfLiteRegistration reg = {};
// This name must match the name specified in the implements signature.
static constexpr char kOpName[] = "my_custom_fused_op";
reg.custom_name = kOpName;
reg.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
// Add your code.
return kTfLiteOk;
};
reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
// Add your code.
return kTfLiteOk;
};
reg.builtin_code = kTfLiteCustom;
resolver->AddCustom(kOpName, &reg);
```
## Converting from composite to fused operation (Advanced)
The overall architecture for converting TensorFlow composite operations to
TensorFlow Lite fused operations is below:
![drawing](../../images/convert/op_fusion.png)
### Wrap the composite operation in a `tf.function`
In the TensorFlow model source code, identify and abstract out the composite
operation into a `tf.function` with the
[experimental\_implements](https://github.com/tensorflow/tensorflow/blob/c11d5d8881fd927165eeb09fd524a80ebaf009f2/tensorflow/python/eager/def_function.py#L470)
function annotation. See an example of [embedding lookup](#composing_ops). The
function defines the interface and its arguments should be used to implement the
conversion logic.
### Write conversion code
The conversion code is written per the interface of the function with the
`implements` annotation. See an example fusion for
[embedding lookup](#fusion_code). Conceptually, the conversion code replaces the
composite implementation of this interface with the fused one.
In the prepare-composite-functions pass, plugin in your
[conversion code](https://github.com/tensorflow/tensorflow/blob/c11d5d8881fd927165eeb09fd524a80ebaf009f2/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L115).
In more advanced usages, it is possible to implement complex transformations of
the composite operation's operands in order to derive the operands of the fused
operation. See
[Keras LSTM](https://github.com/tensorflow/tensorflow/blob/1099faa8d6a941ef44d09ed8c372ff0ffda94112/tensorflow/compiler/mlir/lite/utils/lstm_utils.cc#L627).
conversion code as an example.
### Convert to TensorFlow Lite
Use the
[TFLiteConverter.from_saved_model](https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter#from_saved_model)
API to convert to TensorFlow Lite.
## Under the hood
<a id="under_the_hood"></a>
We now describe high level details of the overall design in converting to fused
operations in TensorFlow Lite.
### Composing operations in TensorFlow
<a id="composing_ops"></a>
The use of `tf.function` with the
[experimental\_implements](https://github.com/tensorflow/tensorflow/blob/c11d5d8881fd927165eeb09fd524a80ebaf009f2/tensorflow/python/eager/def_function.py#L470)
function attribute allows users to explicitly compose new operations using
TensorFlow primitive operations and specify the interface that the resultant
composite operation implements. This is very useful as it provides:
1. A well-defined boundary for the composite operation in the underlying
TensorFlow graph.
1. Explicitly specify the interface that this operation implements. The
arguments of the `tf.function` correspond to the arguments of this
interface.
As an example, lets consider a composite operation defined to implement
embedding lookup. This maps to a fused operation in TensorFlow Lite.
```python
@tf.function(
experimental_implements="embedding_lookup")
def EmbFprop(embs, ids_vec):
"""Embedding forward prop.
Effectively, it computes:
num = size of ids_vec
rets = zeros([num, embedding dim])
for i in range(num):
rets[i, :] = embs[ids_vec[i], :]
return rets
Args:
embs: The embedding matrix.
ids_vec: A vector of int32 embedding ids.
Returns:
The result of embedding lookups. A matrix of shape
[num ids in ids_vec, embedding dims].
"""
num = tf.shape(ids_vec)[0]
rets = inplace_ops.empty([num] + emb_shape_suf, py_utils.FPropDtype(p))
def EmbFpropLoop(i, embs, ids_vec, rets):
# row_id = ids_vec[i]
row_id = tf.gather(ids_vec, i)
# row = embs[row_id]
row = tf.reshape(tf.gather(embs, row_id), [1] + emb_shape_suf)
# rets[i] = row
rets = inplace_ops.alias_inplace_update(rets, [i], row)
return embs, ids_vec, rets
_, _, rets = functional_ops.For(
start=0,
limit=num,
delta=1,
inputs=[embs, ids_vec, rets],
body=EmbFpropLoop,
rewrite_with_while=compiled)
if len(weight_shape) > 2:
rets = tf.reshape(rets, [num, symbolic.ToStatic(p.embedding_dim)])
return rets
```
By making models use composite operations via `tf.function` as illustrated
above, it becomes possible to build a general infrastructure to **identify and
convert** such operations to fused TensorFlow Lite operations.
### Extending the TensorFlow Lite converter
The TensorFlow Lite converter that was released earlier this year only supported
importing TensorFlow models as a graph with all variables replaced with their
corresponding constant values. This does not work for operation fusion since
such graphs have all functions inlined so that the variables can be turned into
constants.
In order to leverage the `tf.function` with the `experimental_implements`
feature during the conversion process, the functions need to be preserved until
later in the conversion process.
As such, we implemented a new workflow of importing and converting TensorFlow
models in the converter to support the composite operation fusion use case.
Specifically, the new features added are:
1. Importing TensorFlow
[saved models into MLIR](https://github.com/tensorflow/tensorflow/blob/1099faa8d6a941ef44d09ed8c372ff0ffda94112/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc#L3748)
1. [fuse composite operations](https://github.com/tensorflow/tensorflow/blob/1099faa8d6a941ef44d09ed8c372ff0ffda94112/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L103)
1. [variable mutability analysis](https://github.com/tensorflow/tensorflow/blob/1099faa8d6a941ef44d09ed8c372ff0ffda94112/tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc#L43)
1. [freeze all read-only variables](https://github.com/tensorflow/tensorflow/blob/1099faa8d6a941ef44d09ed8c372ff0ffda94112/tensorflow/compiler/mlir/tensorflow/transforms/freeze_global_tensors.cc#L44)
This allows us to perform operation fusion using the functions representing the
composite operations prior to function inlining and variable freezing.
### Implementing operation fusion
Lets look at the operation fusion pass in more detail. This pass does the
following:
1. Loop through all functions in the MLIR module.
1. If a function has the tf.\_implements attribute, based on the attribute
value, calls the appropriate operation fusion utility.
1. The operation fusion utility operates on the functions operands and
attributes (which serve as the interface for the conversion) and replaces
the body of the function with an equivalent function body containing the
fused operation.
1. In many cases, the replaced body will contain operations other than the
fused operation. These correspond to some static transforms on the
functions operands in order to obtain the operands of the fused operation.
Since these computations can all be constant folded away, they would not be
present in the exported flatbuffer where only the fused operation would
exist.
Here is code snippet from the pass showing the main workflow:
```
void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
StringAttr attr) {
if (attr.getValue() == "embedding_lookup") {
func.eraseBody();
func.addEntryBlock();
// Convert the composite embedding_lookup function body to a
// TFLite fused embedding_lookup op.
ConvertEmbeddedLookupFunc convert_embedded_lookup(func);
if (failed(convert_embedded_lookup.VerifySignature())) {
return signalPassFailure();
}
convert_embedded_lookup.RewriteFunc();
} else if (attr.getValue() == mlir::TFL::kKerasLstm) {
func.eraseBody();
func.addEntryBlock();
OpBuilder builder(func.getBody());
if (failed(ConvertKerasLSTMLayer(func, &builder))) {
return signalPassFailure();
}
} else if (.....) /* Other fusions can plug in here */
}
```
Here is code snippet showing mapping this composite operation to a fused
operation in TensorFlow Lite leveraging the function as a conversion interface.
<a id="fusion_code"></a>
```c++
void RewriteFunc() {
Value lookup = func_.getArgument(1);
Value value = func_.getArgument(0);
auto output_type = func_.getType().getResult(0);
OpBuilder builder(func_.getBody());
auto op = builder.create<mlir::TFL::EmbeddingLookupOp>(
func_.getLoc(), output_type, lookup, value);
builder.create<mlir::ReturnOp>(func_.getLoc(), op.getResult());
}
```
+198
View File
@@ -0,0 +1,198 @@
# TensorFlow RNN conversion to TensorFlow Lite
## Overview
TensorFlow Lite supports converting TensorFlow RNN models to TensorFlow Lites
fused LSTM operations. Fused operations exist to maximize the performance of
their underlying kernel implementations, as well as provide a higher level
interface to define complex transformations like quantizatization.
Since there are many variants of RNN APIs in TensorFlow, our approach has been
two fold:
1. Provide **native support for standard TensorFlow RNN APIs** like Keras LSTM.
This is the recommended option.
1. Provide an **interface** **into the conversion infrastructure for**
**user-defined** **RNN implementations** to plug in and get converted to
TensorFlow Lite. We provide a couple of out of box examples of such
conversion using lingvos
[LSTMCellSimple](https://github.com/tensorflow/tensorflow/blob/82abf0dbf316526cd718ae8cd7b11cfcb805805e/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L130)
and
[LayerNormalizedLSTMCellSimple](https://github.com/tensorflow/tensorflow/blob/c11d5d8881fd927165eeb09fd524a80ebaf009f2/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L137)
RNN interfaces.
## Converter API
The feature is part of TensorFlow 2.3 release. It is also available through the
[tf-nightly](https://pypi.org/project/tf-nightly/) pip or from head.
This conversion functionality is available when converting to TensorFlow Lite
via a SavedModel or from the Keras model directly. See example usages.
### From saved model
<a id="from_saved_model"></a>
```
# build a saved model. Here concrete_function is the exported function
# corresponding to the TensorFlow model containing one or more
# Keras LSTM layers.
saved_model, saved_model_dir = build_saved_model_lstm(...)
saved_model.save(saved_model_dir, save_format="tf", signatures=concrete_func)
# Convert the model.
converter = TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()
```
### From Keras model
```
# build a Keras model
keras_model = build_keras_lstm(...)
# Convert the model.
converter = TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()
```
## Example
Keras LSTM to TensorFlow Lite
[Colab](https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/experimental_new_converter/Keras_LSTM_fusion_Codelab.ipynb)
illustrates the end to end usage with the TensorFlow Lite interpreter.
## TensorFlow RNNs APIs supported
<a id="rnn_apis"></a>
### Keras LSTM conversion (recommended)
We support out-of-the-box conversion of Keras LSTM to TensorFlow Lite. For
details on how this works please refer to the
[Keras LSTM interface](https://github.com/tensorflow/tensorflow/blob/35a3ab91b42503776f428bda574b74b9a99cd110/tensorflow/python/keras/layers/recurrent_v2.py#L1238)<span style="text-decoration:space;">
</span>and to the conversion logic
[here](https://github.com/tensorflow/tensorflow/blob/35a3ab91b42503776f428bda574b74b9a99cd110/tensorflow/compiler/mlir/lite/utils/lstm_utils.cc#L627).
Also important is to highlight the TensorFlow Lites LSTM contract with respect
to the Keras operation definition:
1. The dimension 0 of the **input** tensor is the batch size.
1. The dimension 0 of the **recurrent\_weight** tensor is the number of
outputs.
1. The **weight** and **recurrent\_kernel** tensors are transposed.
1. The transposed weight, transposed recurrent\_kernel and **bias** tensors are
split into 4 equal sized tensors along the dimension 0. These correspond to
**input gate, forget gate, cell, and output gate**.
#### Keras LSTM Variants
##### Time major
Users may choose time-major or no time-major. Keras LSTM adds a time-major
attribute in the function def attributes. For Unidirectional sequence LSTM, we
can simply map to unidirecional\_sequence\_lstm's
[time major attribute](https://github.com/tensorflow/tensorflow/blob/35a3ab91b42503776f428bda574b74b9a99cd110/tensorflow/compiler/mlir/lite/ir/tfl_ops.td#L3902).
##### BiDirectional LSTM
Bidirectional LSTM can be implemented with two Keras LSTM layers, one for
forward and one for backward, see examples
[here](https://github.com/tensorflow/tensorflow/blob/35a3ab91b42503776f428bda574b74b9a99cd110/tensorflow/python/keras/layers/wrappers.py#L382).
Once we see the go\_backward attribute, we recognize it as backward LSTM, then
we group forward & backward LSTM together. **This is future work.** Currently,
this creates two UnidirectionalSequenceLSTM operations in the TensorFlow Lite
model.
### User-defined LSTM conversion examples
TensorFlow Lite also provides a way to convert user defined LSTM
implementations. Here we use Lingvos LSTM as an example of how that can be
implemented. For details please refer to the
[lingvo.LSTMCellSimple interface](https://github.com/tensorflow/lingvo/blob/91a4609dbc2579748a95110eda59c66d17c594c5/lingvo/core/rnn_cell.py#L228)
and the conversion logic
[here](https://github.com/tensorflow/tensorflow/blob/82abf0dbf316526cd718ae8cd7b11cfcb805805e/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L130).
We also provide an example for another of Lingvos LSTM definitions in
[lingvo.LayerNormalizedLSTMCellSimple interface](https://github.com/tensorflow/lingvo/blob/91a4609dbc2579748a95110eda59c66d17c594c5/lingvo/core/rnn_cell.py#L1173)
and its conversion logic
[here](https://github.com/tensorflow/tensorflow/blob/c11d5d8881fd927165eeb09fd524a80ebaf009f2/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L137).
## “Bring your own TensorFlow RNN” to TensorFlow Lite
If a user's RNN interface is different from the standard supported ones, there
are a couple of options:
**Option 1:** Write adapter code in TensorFlow python to adapt the RNN interface
to the Keras RNN interface. This means a tf.function with
[tf\_implements annotation](https://github.com/tensorflow/community/pull/113) on
the generated RNN interfaces function that is identical to the one generated by
the Keras LSTM layer. After this, the same conversion API used for Keras LSTM
will work.
**Option 2:** If the above is not possible (e.g. the Keras LSTM is missing some
functionality that is currently exposed by TensorFlow Lites fused LSTM op like
layer normalization), then extend the TensorFlow Lite converter by writing
custom conversion code and plug it into the prepare-composite-functions
MLIR-pass
[here](https://github.com/tensorflow/tensorflow/blob/c11d5d8881fd927165eeb09fd524a80ebaf009f2/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L115).
The functions interface should be treated like an API contract and should
contain the arguments needed to convert to fused TensorFlow Lite LSTM
operations - i.e. input, bias, weights, projection, layer normalization, etc. It
is preferable for the tensors passed as arguments to this function to have known
rank (i.e. RankedTensorType in MLIR). This makes it much easier to write
conversion code that can assume these tensors as RankedTensorType and helps
transform them to ranked tensors corresponding to the fused TensorFlow Lite
operators operands.
A complete example of such conversion flow is Lingvos LSTMCellSimple to
TensorFlow Lite conversion.
The LSTMCellSimple in Lingvo is defined
[here](https://github.com/tensorflow/lingvo/blob/91a4609dbc2579748a95110eda59c66d17c594c5/lingvo/core/rnn_cell.py#L228).
Models trained with this LSTM cell can be converted to TensorFlow Lite as
follows:
1. Wrap all uses of LSTMCellSimple in a tf.function with a tf\_implements
annotation that is labelled as such (e.g. lingvo.LSTMCellSimple would be a
good annotation name here). Make sure the tf.function that is generated
matches the interface of the function expected in the conversion code. This
is a contract between the model author adding the annotation and the
conversion code.
1. Extend the prepare-composite-functions pass to plug in a custom composite op
to TensorFlow Lite fused LSTM op conversion. See
[LSTMCellSimple](https://github.com/tensorflow/tensorflow/blob/82abf0dbf316526cd718ae8cd7b11cfcb805805e/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L130)
conversion code.
The conversion contract:
1. **Weight** and **projection** tensors are transposed.
1. The **{input, recurrent}** to **{cell, input gate, forget gate, output
gate}** are extracted by slicing the transposed weight tensor.
1. The **{bias}** to **{cell, input gate, forget gate, output gate}** are
extracted by slicing the bias tensor.
1. The **projection** is extracted by slicing the transposed projection tensor.
1. Similar conversion is written for
[LayerNormalizedLSTMCellSimple](https://github.com/tensorflow/tensorflow/blob/c11d5d8881fd927165eeb09fd524a80ebaf009f2/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L137).
1. The rest of the TensorFlow Lite conversion infrastructure, including all the
[MLIR passes](https://github.com/tensorflow/tensorflow/blob/35a3ab91b42503776f428bda574b74b9a99cd110/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc#L57)
defined as well as the final export to TensorFlow Lite flatbuffer can be
reused.
## Known issues/limitations
1. Currently there is support only for converting stateless Keras LSTM (default
behavior in Keras). Stateful Keras LSTM conversion is future work.
1. It is still possible to model a stateful Keras LSTM layer using the
underlying stateless Keras LSTM layer and managing the state explicitly in
the user program. Such a TensorFlow program can still be converted to
TensorFlow Lite using the feature being described here.
1. Bidirectional LSTM is currently modelled as two UnidirectionalSequenceLSTM
operations in TensorFlow Lite. This will be replaced with a single
BidirectionalSequenceLSTM op.
@@ -0,0 +1,646 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "3XX46cTrh6iD"
},
"source": [
"##### Copyright 2024 The AI Edge Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "sKrlWr6Kh-mF"
},
"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": "hST65kOHXpiL"
},
"source": [
"# Transfer Learning for the Audio Domain with TensorFlow Lite Model Maker\n",
"\n",
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://www.tensorflow.org/lite/models/modify/model_maker/audio_classification\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/audio_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/audio_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/models/modify/model_maker/audio_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://tfhub.dev/google/yamnet/1\"><img src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" />See TF Hub model</a>\n",
" </td>\n",
"\n",
"</table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BB5k6xNKJ5Xe"
},
"source": [
"\n",
"In this colab notebook, you'll learn how to use the [TensorFlow Lite Model Maker](https://www.tensorflow.org/lite/models/modify/model_maker) to train a custom audio classification model.\n",
"\n",
"The Model Maker library uses transfer learning to simplify the process of training a TensorFlow Lite model using a custom dataset. Retraining a TensorFlow Lite model with your own custom dataset reduces the amount of training data and time required.\n",
"\n",
"It is part of the [Codelab to Customize an Audio model and deploy on Android](https://codelabs.developers.google.com/codelabs/tflite-audio-classification-custom-model-android).\n",
"\n",
"You'll use a custom birds dataset and export a TFLite model that can be used on a phone, a TensorFlow.JS model that can be used for inference in the browser and also a SavedModel version that you can use for serving.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AeZZ_cSsZfPx"
},
"source": [
"## Intalling dependencies\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "wbMc4vHjaYdQ"
},
"outputs": [],
"source": [
"!sudo apt -y install libportaudio2\n",
"!pip install tflite-model-maker"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "z2ck_Ghdcgt9"
},
"source": [
"## Import TensorFlow, Model Maker and other libraries\n",
"\n",
"Among the dependencies that are needed, you'll use TensorFlow and Model Maker. Aside those, the others are for audio manipulation, playing and visualizations."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rwUA9u4oWoCR"
},
"outputs": [],
"source": [
"import tensorflow as tf\n",
"import tflite_model_maker as mm\n",
"from tflite_model_maker import audio_classifier\n",
"import os\n",
"\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"import seaborn as sns\n",
"\n",
"import itertools\n",
"import glob\n",
"import random\n",
"\n",
"from IPython.display import Audio, Image\n",
"from scipy.io import wavfile\n",
"\n",
"print(f\"TensorFlow Version: {tf.__version__}\")\n",
"print(f\"Model Maker Version: {mm.__version__}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HIfm2TxKZAuA"
},
"source": [
"## The Birds dataset\n",
"\n",
"The Birds dataset is an education collection of 5 types of birds songs:\n",
"\n",
"- White-breasted Wood-Wren\n",
"- House Sparrow\n",
"- Red Crossbill\n",
"- Chestnut-crowned Antpitta\n",
"- Azara's Spinetail\n",
"\n",
"The original audio came from [Xeno-canto](https://www.xeno-canto.org/) which is a website dedicated to sharing bird sounds from all over the world.\n",
"\n",
"Let's start by downloading the data."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "upNRfilkNSmr"
},
"outputs": [],
"source": [
"birds_dataset_folder = tf.keras.utils.get_file('birds_dataset.zip',\n",
" 'https://storage.googleapis.com/laurencemoroney-blog.appspot.com/birds_dataset.zip',\n",
" cache_dir='./',\n",
" cache_subdir='dataset',\n",
" extract=True)\n",
" "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "441bbzZ5d6oq"
},
"source": [
"## Explore the data\n",
"\n",
"The audios are already split in train and test folders. Inside each split folder, there's one folder for each bird, using their `bird_code` as name.\n",
"\n",
"The audios are all mono and with 16kHz sample rate.\n",
"\n",
"For more information about each file, you can read the `metadata.csv` file. It contains all the files authors, lincenses and some more information. You won't need to read it yourself on this tutorial."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ayd7UqCfQQFU"
},
"outputs": [],
"source": [
"# @title [Run this] Util functions and data structures.\n",
"\n",
"data_dir = './dataset/small_birds_dataset'\n",
"\n",
"bird_code_to_name = {\n",
" 'wbwwre1': 'White-breasted Wood-Wren',\n",
" 'houspa': 'House Sparrow',\n",
" 'redcro': 'Red Crossbill', \n",
" 'chcant2': 'Chestnut-crowned Antpitta',\n",
" 'azaspi1': \"Azara's Spinetail\", \n",
"}\n",
"\n",
"birds_images = {\n",
" 'wbwwre1': 'https://upload.wikimedia.org/wikipedia/commons/thumb/2/22/Henicorhina_leucosticta_%28Cucarachero_pechiblanco%29_-_Juvenil_%2814037225664%29.jpg/640px-Henicorhina_leucosticta_%28Cucarachero_pechiblanco%29_-_Juvenil_%2814037225664%29.jpg', # \tAlejandro Bayer Tamayo from Armenia, Colombia \n",
" 'houspa': 'https://upload.wikimedia.org/wikipedia/commons/thumb/5/52/House_Sparrow%2C_England_-_May_09.jpg/571px-House_Sparrow%2C_England_-_May_09.jpg', # \tDiliff\n",
" 'redcro': 'https://upload.wikimedia.org/wikipedia/commons/thumb/4/49/Red_Crossbills_%28Male%29.jpg/640px-Red_Crossbills_%28Male%29.jpg', # Elaine R. Wilson, www.naturespicsonline.com\n",
" 'chcant2': 'https://upload.wikimedia.org/wikipedia/commons/thumb/6/67/Chestnut-crowned_antpitta_%2846933264335%29.jpg/640px-Chestnut-crowned_antpitta_%2846933264335%29.jpg', # \tMike's Birds from Riverside, CA, US\n",
" 'azaspi1': 'https://upload.wikimedia.org/wikipedia/commons/thumb/b/b2/Synallaxis_azarae_76608368.jpg/640px-Synallaxis_azarae_76608368.jpg', # https://www.inaturalist.org/photos/76608368\n",
"}\n",
"\n",
"test_files = os.path.abspath(os.path.join(data_dir, 'test/*/*.wav'))\n",
"\n",
"def get_random_audio_file():\n",
" test_list = glob.glob(test_files)\n",
" random_audio_path = random.choice(test_list)\n",
" return random_audio_path\n",
"\n",
"\n",
"def show_bird_data(audio_path):\n",
" sample_rate, audio_data = wavfile.read(audio_path, 'rb')\n",
"\n",
" bird_code = audio_path.split('/')[-2]\n",
" print(f'Bird name: {bird_code_to_name[bird_code]}')\n",
" print(f'Bird code: {bird_code}')\n",
" display(Image(birds_images[bird_code]))\n",
"\n",
" plttitle = f'{bird_code_to_name[bird_code]} ({bird_code})'\n",
" plt.title(plttitle)\n",
" plt.plot(audio_data)\n",
" display(Audio(audio_data, rate=sample_rate))\n",
"\n",
"print('functions and data structures created')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Yrv0uD7aXYl4"
},
"source": [
"### Playing some audio\n",
"\n",
"To have a better understanding about the data, lets listen to a random audio files from the test split.\n",
"\n",
"Note: later in this notebook you'll run inference on this audio for testing"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tEeMZh-VQy97"
},
"outputs": [],
"source": [
"random_audio = get_random_audio_file()\n",
"show_bird_data(random_audio)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pQj1Mf7YZELS"
},
"source": [
"## Training the Model\n",
"\n",
"When using Model Maker for audio, you have to start with a model spec. This is the base model that your new model will extract information to learn about the new classes. It also affects how the dataset will be transformed to respect the models spec parameters like: sample rate, number of channels.\n",
"\n",
"[YAMNet](https://tfhub.dev/google/yamnet/1) is an audio event classifier trained on the AudioSet dataset to predict audio events from the AudioSet ontology.\n",
"\n",
"It's input is expected to be at 16kHz and with 1 channel.\n",
"\n",
"You don't need to do any resampling yourself. Model Maker takes care of that for you.\n",
"\n",
"- `frame_length` is to decide how long each traininng sample is. in this caase EXPECTED_WAVEFORM_LENGTH * 3s\n",
"\n",
"- `frame_steps` is to decide how far appart are the training samples. In this case, the ith sample will start at EXPECTED_WAVEFORM_LENGTH * 6s after the (i-1)th sample.\n",
"\n",
"The reason to set these values is to work around some limitation in real world dataset.\n",
"\n",
"For example, in the bird dataset, birds don't sing all the time. They sing, rest and sing again, with noises in between. Having a long frame would help capture the singing, but setting it too long will reduce the number of samples for training.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tUcxtfHXY7XS"
},
"outputs": [],
"source": [
"spec = audio_classifier.YamNetSpec(\n",
" keep_yamnet_and_custom_heads=True,\n",
" frame_step=3 * audio_classifier.YamNetSpec.EXPECTED_WAVEFORM_LENGTH,\n",
" frame_length=6 * audio_classifier.YamNetSpec.EXPECTED_WAVEFORM_LENGTH)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EF185yZ_M7zu"
},
"source": [
"## Loading the data\n",
"\n",
"Model Maker has the API to load the data from a folder and have it in the expected format for the model spec.\n",
"\n",
"The train and test split are based on the folders. The validation dataset will be created as 20% of the train split.\n",
"\n",
"Note: The `cache=True` is important to make training later faster but it will also require more RAM to hold the data. For the birds dataset that is not a problem since it's only 300MB, but if you use your own data you have to pay attention to it.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cX0RqETqZgzo"
},
"outputs": [],
"source": [
"train_data = audio_classifier.DataLoader.from_folder(\n",
" spec, os.path.join(data_dir, 'train'), cache=True)\n",
"train_data, validation_data = train_data.split(0.8)\n",
"test_data = audio_classifier.DataLoader.from_folder(\n",
" spec, os.path.join(data_dir, 'test'), cache=True)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ziMghju-Rts2"
},
"source": [
"## Training the model\n",
"\n",
"the audio_classifier has the [`create`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/audio_classifier/create) method that creates a model and already start training it. \n",
"\n",
"You can customize many parameterss, for more information you can read more details in the documentation.\n",
"\n",
"On this first try you'll use all the default configurations and train for 100 epochs.\n",
"\n",
"Note: The first epoch takes longer than all the other ones because it's when the cache is created. After that each epoch takes close to 1 second."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8r6Awvl4ZkIv"
},
"outputs": [],
"source": [
"batch_size = 128\n",
"epochs = 100\n",
"\n",
"print('Training the model')\n",
"model = audio_classifier.create(\n",
" train_data,\n",
" spec,\n",
" validation_data,\n",
" batch_size=batch_size,\n",
" epochs=epochs)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "oXMEHZkAxJTl"
},
"source": [
"The accuracy looks good but it's important to run the evaluation step on the test data and vefify your model achieved good results on unseed data."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xEM1aRNKtvHS"
},
"source": [
""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "GDoQACMrZnOx"
},
"outputs": [],
"source": [
"print('Evaluating the model')\n",
"model.evaluate(test_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8QRRAM39aOxS"
},
"source": [
"## Understanding your model\n",
"\n",
"When training a classifier, it's useful to see the [confusion matrix](https://en.wikipedia.org/wiki/Confusion_matrix). The confusion matrix gives you detailed knowledge of how your classifier is performing on test data.\n",
"\n",
"Model Maker already creates the confusion matrix for you."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "zqB3c0368iH3"
},
"outputs": [],
"source": [
"def show_confusion_matrix(confusion, test_labels):\n",
" \"\"\"Compute confusion matrix and normalize.\"\"\"\n",
" confusion_normalized = confusion.astype(\"float\") / confusion.sum(axis=1)\n",
" axis_labels = test_labels\n",
" ax = sns.heatmap(\n",
" confusion_normalized, xticklabels=axis_labels, yticklabels=axis_labels,\n",
" cmap='Blues', annot=True, fmt='.2f', square=True)\n",
" plt.title(\"Confusion matrix\")\n",
" plt.ylabel(\"True label\")\n",
" plt.xlabel(\"Predicted label\")\n",
"\n",
"confusion_matrix = model.confusion_matrix(test_data)\n",
"show_confusion_matrix(confusion_matrix.numpy(), test_data.index_to_label)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7gr1s7juBy7H"
},
"source": [
"## Testing the model [Optional]\n",
"\n",
"You can try the model on a sample audio from the test dataset just to see the results.\n",
"\n",
"First you get the serving model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "PmlmTl42Bq_u"
},
"outputs": [],
"source": [
"serving_model = model.create_serving_model()\n",
"\n",
"print(f'Model\\'s input shape and type: {serving_model.inputs}')\n",
"print(f'Model\\'s output shape and type: {serving_model.outputs}')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "eQsZFO2mrYhx"
},
"source": [
"Coming back to the random audio you loaded earlier"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8dv5ViK0reXc"
},
"outputs": [],
"source": [
"# if you want to try another file just uncoment the line below\n",
"random_audio = get_random_audio_file()\n",
"show_bird_data(random_audio)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uixOfKSUj_9m"
},
"source": [
"The model created has a fixed input window. \n",
"\n",
"For a given audio file, you'll have to split it in windows of data of the expected size. The last window might need to be filled with zeros."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "YAvGKQL0lNty"
},
"outputs": [],
"source": [
"sample_rate, audio_data = wavfile.read(random_audio, 'rb')\n",
"\n",
"audio_data = np.array(audio_data) / tf.int16.max\n",
"input_size = serving_model.input_shape[1]\n",
"\n",
"splitted_audio_data = tf.signal.frame(audio_data, input_size, input_size, pad_end=True, pad_value=0)\n",
"\n",
"print(f'Test audio path: {random_audio}')\n",
"print(f'Original size of the audio data: {len(audio_data)}')\n",
"print(f'Number of windows for inference: {len(splitted_audio_data)}')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "PLxKd0eFkMcR"
},
"source": [
"You'll loop over all the splitted audio and apply the model for each one of them.\n",
"\n",
"The model you've just trained has 2 outputs: The original YAMNet's output and the one you've just trained. This is important because the real world environment is more complicated than just bird sounds. You can use the YAMNet's output to filter out non relevant audio, for example, on the birds use case, if YAMNet is not classifying Birds or Animals, this might show that the output from your model might have an irrelevant classification.\n",
"\n",
"\n",
"Below both outpus are printed to make it easier to understand their relation. Most of the mistakes that your model make are when YAMNet's prediction is not related to your domain (eg: birds)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4-8fJLrxGwYT"
},
"outputs": [],
"source": [
"print(random_audio)\n",
"\n",
"results = []\n",
"print('Result of the window ith: your model class -> score, (spec class -> score)')\n",
"for i, data in enumerate(splitted_audio_data):\n",
" yamnet_output, inference = serving_model(data)\n",
" results.append(inference[0].numpy())\n",
" result_index = tf.argmax(inference[0])\n",
" spec_result_index = tf.argmax(yamnet_output[0])\n",
" t = spec._yamnet_labels()[spec_result_index]\n",
" result_str = f'Result of the window {i}: ' \\\n",
" f'\\t{test_data.index_to_label[result_index]} -> {inference[0][result_index].numpy():.3f}, ' \\\n",
" f'\\t({spec._yamnet_labels()[spec_result_index]} -> {yamnet_output[0][spec_result_index]:.3f})'\n",
" print(result_str)\n",
"\n",
"\n",
"results_np = np.array(results)\n",
"mean_results = results_np.mean(axis=0)\n",
"result_index = mean_results.argmax()\n",
"print(f'Mean result: {test_data.index_to_label[result_index]} -> {mean_results[result_index]}')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yASrikBgZ9ZO"
},
"source": [
"## Exporting the model\n",
"\n",
"The last step is exporting your model to be used on embedded devices or on the browser.\n",
"\n",
"The `export` method export both formats for you."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Xw_ehPxAdQlz"
},
"outputs": [],
"source": [
"models_path = './birds_models'\n",
"print(f'Exporing the TFLite model to {models_path}')\n",
"\n",
"model.export(models_path, tflite_filename='my_birds_model.tflite')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "jjZRKmurA3y_"
},
"source": [
"You can also export the SavedModel version for serving or using on a Python environment."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "veBwppOsA-kn"
},
"outputs": [],
"source": [
"model.export(models_path, export_format=[mm.ExportFormat.SAVED_MODEL, mm.ExportFormat.LABEL])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5xr0idac6xfi"
},
"source": [
"## Next Steps\n",
"\n",
"You did it.\n",
"\n",
"Now your new model can be deployed on mobile devices using [TFLite AudioClassifier Task API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/audio_classifier).\n",
"\n",
"You can also try the same process with your own data with different classes and here is the documentation for [Model Maker for Audio Classification](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/audio_classifier).\n",
"\n",
"Also learn from end-to-end reference apps: [Android](https://github.com/tensorflow/examples/tree/master/lite/examples/sound_classification/android/), [iOS](https://github.com/tensorflow/examples/tree/master/lite/examples/sound_classification/ios)."
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "model_maker_audio_classification.ipynb",
"private_outputs": true,
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
@@ -0,0 +1,925 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "h2q27gKz1H20"
},
"source": [
"##### Copyright 2024 The AI Edge Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "TUfAcER1oUS6"
},
"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": "Gb7qyhNL1yWt"
},
"source": [
"# Image classification with TensorFlow Lite Model Maker"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "nDABAblytltI"
},
"source": [
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://www.tensorflow.org/lite/models/modify/model_maker/image_classification\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/image_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/image_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/models/modify/model_maker/image_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1\"><img src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" />See TF Hub model</a>\n",
" </td>\n",
"</table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "m86-Nh4pMHqY"
},
"source": [
"The [TensorFlow Lite Model Maker library](https://www.tensorflow.org/lite/models/modify/model_maker) simplifies the process of adapting and converting a TensorFlow neural-network model to particular input data when deploying this model for on-device ML applications.\n",
"\n",
"This notebook shows an end-to-end example that utilizes this Model Maker library to illustrate the adaption and conversion of a commonly-used image classification model to classify flowers on a mobile device."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bcLF2PKkSbV3"
},
"source": [
"## Prerequisites\n",
"\n",
"To run this example, we first need to install several required packages, including Model Maker package that in GitHub [repo](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6cv3K3oaksJv"
},
"outputs": [],
"source": [
"!sudo apt -y install libportaudio2\n",
"!pip install -q tflite-model-maker"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Gx1HGRoFQ54j"
},
"source": [
"Import the required packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "XtxiUeZEiXpt"
},
"outputs": [],
"source": [
"import os\n",
"\n",
"import numpy as np\n",
"\n",
"import tensorflow as tf\n",
"assert tf.__version__.startswith('2')\n",
"\n",
"from tflite_model_maker import model_spec\n",
"from tflite_model_maker import image_classifier\n",
"from tflite_model_maker.config import ExportFormat\n",
"from tflite_model_maker.config import QuantizationConfig\n",
"from tflite_model_maker.image_classifier import DataLoader\n",
"\n",
"import matplotlib.pyplot as plt"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "KKRaYHABpob5"
},
"source": [
"## Simple End-to-End Example"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "SiZZ5DHXotaW"
},
"source": [
"### Get the data path\n",
"\n",
"Let's get some images to play with this simple end-to-end example. Hundreds of images is a good start for Model Maker while more data could achieve better accuracy."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "3jz5x0JoskPv"
},
"outputs": [],
"source": [
"image_path = tf.keras.utils.get_file(\n",
" 'flower_photos.tgz',\n",
" 'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',\n",
" extract=True)\n",
"image_path = os.path.join(os.path.dirname(image_path), 'flower_photos')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "a55MR6i6nuDm"
},
"source": [
"You could replace `image_path` with your own image folders. As for uploading data to colab, you could find the upload button in the left sidebar shown in the image below with the red rectangle. Just have a try to upload a zip file and unzip it. The root file path is the current path.\n",
"\n",
"<img src=\"https://storage.googleapis.com/download.tensorflow.org/models/tflite/screenshots/model_maker_image_classification.png\" alt=\"Upload File\" width=\"800\" hspace=\"100\">"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "NNRNv_mloS89"
},
"source": [
"If you prefer not to upload your images to the cloud, you could try to run the library locally following the [guide](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker) in GitHub."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "w-VDriAdsowu"
},
"source": [
"### Run the example\n",
"The example just consists of 4 lines of code as shown below, each of which representing one step of the overall process.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6ahtcO86tZBL"
},
"source": [
"Step 1. Load input data specific to an on-device ML app. Split it into training data and testing data."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "lANoNS_gtdH1"
},
"outputs": [],
"source": [
"data = DataLoader.from_folder(image_path)\n",
"train_data, test_data = data.split(0.9)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Y_9IWyIztuRF"
},
"source": [
"Step 2. Customize the TensorFlow model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yRXMZbrwtyRD"
},
"outputs": [],
"source": [
"model = image_classifier.create(train_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "oxU2fDr-t2Ya"
},
"source": [
"Step 3. Evaluate the model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "wQr02VxJt6Cs"
},
"outputs": [],
"source": [
"loss, accuracy = model.evaluate(test_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "eVZw9zU8t84y"
},
"source": [
"Step 4. Export to TensorFlow Lite model.\n",
"\n",
"Here, we export TensorFlow Lite model with [metadata](https://www.tensorflow.org/lite/models/convert/metadata) which provides a standard for model descriptions. The label file is embedded in metadata. The default post-training quantization technique is full integer quantization for the image classification task.\n",
"\n",
"You could download it in the left sidebar same as the uploading part for your own use."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Zb-eIzfluCoa"
},
"outputs": [],
"source": [
"model.export(export_dir='.')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pyju1qc_v-wy"
},
"source": [
"After these simple 4 steps, we could further use TensorFlow Lite model file in on-device applications like in [image classification](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification) reference app."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "R1QG32ivs9lF"
},
"source": [
"## Detailed Process\n",
"\n",
"Currently, we support several models such as EfficientNet-Lite* models, MobileNetV2, ResNet50 as pre-trained models for image classification. But it is very flexible to add new pre-trained models to this library with just a few lines of code.\n",
"\n",
"\n",
"The following walks through this end-to-end example step by step to show more detail."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ygEncJxtl-nQ"
},
"source": [
"### Step 1: Load Input Data Specific to an On-device ML App\n",
"\n",
"The flower dataset contains 3670 images belonging to 5 classes. Download the archive version of the dataset and untar it.\n",
"\n",
"The dataset has the following directory structure:\n",
"\n",
"<pre>\n",
"<b>flower_photos</b>\n",
"|__ <b>daisy</b>\n",
" |______ 100080576_f52e8ee070_n.jpg\n",
" |______ 14167534527_781ceb1b7a_n.jpg\n",
" |______ ...\n",
"|__ <b>dandelion</b>\n",
" |______ 10043234166_e6dd915111_n.jpg\n",
" |______ 1426682852_e62169221f_m.jpg\n",
" |______ ...\n",
"|__ <b>roses</b>\n",
" |______ 102501987_3cdb8e5394_n.jpg\n",
" |______ 14982802401_a3dfb22afb.jpg\n",
" |______ ...\n",
"|__ <b>sunflowers</b>\n",
" |______ 12471791574_bb1be83df4.jpg\n",
" |______ 15122112402_cafa41934f.jpg\n",
" |______ ...\n",
"|__ <b>tulips</b>\n",
" |______ 13976522214_ccec508fe7.jpg\n",
" |______ 14487943607_651e8062a1_m.jpg\n",
" |______ ...\n",
"</pre>"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7tOfUr2KlgpU"
},
"outputs": [],
"source": [
"image_path = tf.keras.utils.get_file(\n",
" 'flower_photos.tgz',\n",
" 'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',\n",
" extract=True)\n",
"image_path = os.path.join(os.path.dirname(image_path), 'flower_photos')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "E051HBUM5owi"
},
"source": [
"Use `DataLoader` class to load data.\n",
"\n",
"As for `from_folder()` method, it could load data from the folder. It assumes that the image data of the same class are in the same subdirectory and the subfolder name is the class name. Currently, JPEG-encoded images and PNG-encoded images are supported."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "I_fOlZsklmlL"
},
"outputs": [],
"source": [
"data = DataLoader.from_folder(image_path)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "u501eT4koURB"
},
"source": [
"Split it to training data (80%), validation data (10%, optional) and testing data (10%)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cY4UU5SUobtJ"
},
"outputs": [],
"source": [
"train_data, rest_data = data.split(0.8)\n",
"validation_data, test_data = rest_data.split(0.5)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Z9_MYPie3EMO"
},
"source": [
"Show 25 image examples with labels."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Ih4Wx44I482b"
},
"outputs": [],
"source": [
"plt.figure(figsize=(10,10))\n",
"for i, (image, label) in enumerate(data.gen_dataset().unbatch().take(25)):\n",
" plt.subplot(5,5,i+1)\n",
" plt.xticks([])\n",
" plt.yticks([])\n",
" plt.grid(False)\n",
" plt.imshow(image.numpy(), cmap=plt.cm.gray)\n",
" plt.xlabel(data.index_to_label[label.numpy()])\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AWuoensX4vDA"
},
"source": [
"### Step 2: Customize the TensorFlow Model\n",
"\n",
"Create a custom image classifier model based on the loaded data. The default model is EfficientNet-Lite0.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "TvYSUuJY3QxR"
},
"outputs": [],
"source": [
"model = image_classifier.create(train_data, validation_data=validation_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "4JFOKWnH9x8_"
},
"source": [
"Have a look at the detailed model structure."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "QNXAfjl192dC"
},
"outputs": [],
"source": [
"model.summary()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LP5FPk_tOxoZ"
},
"source": [
"### Step 3: Evaluate the Customized Model\n",
"\n",
"Evaluate the result of the model, get the loss and accuracy of the model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "A8c2ZQ0J3Riy"
},
"outputs": [],
"source": [
"loss, accuracy = model.evaluate(test_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6ZCrYOWoCt05"
},
"source": [
"We could plot the predicted results in 100 test images. Predicted labels with red color are the wrong predicted results while others are correct."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "n9O9Kx7nDQWD"
},
"outputs": [],
"source": [
"# A helper function that returns 'red'/'black' depending on if its two input\n",
"# parameter matches or not.\n",
"def get_label_color(val1, val2):\n",
" if val1 == val2:\n",
" return 'black'\n",
" else:\n",
" return 'red'\n",
"\n",
"# Then plot 100 test images and their predicted labels.\n",
"# If a prediction result is different from the label provided label in \"test\"\n",
"# dataset, we will highlight it in red color.\n",
"plt.figure(figsize=(20, 20))\n",
"predicts = model.predict_top_k(test_data)\n",
"for i, (image, label) in enumerate(test_data.gen_dataset().unbatch().take(100)):\n",
" ax = plt.subplot(10, 10, i+1)\n",
" plt.xticks([])\n",
" plt.yticks([])\n",
" plt.grid(False)\n",
" plt.imshow(image.numpy(), cmap=plt.cm.gray)\n",
"\n",
" predict_label = predicts[i][0][0]\n",
" color = get_label_color(predict_label,\n",
" test_data.index_to_label[label.numpy()])\n",
" ax.xaxis.label.set_color(color)\n",
" plt.xlabel('Predicted: %s' % predict_label)\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "S3H0rkbLUZAG"
},
"source": [
"If the accuracy doesn't meet the app requirement, one could refer to [Advanced Usage](#scrollTo=zNDBP2qA54aK) to explore alternatives such as changing to a larger model, adjusting re-training parameters etc."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aeHoGAceO2xV"
},
"source": [
"### Step 4: Export to TensorFlow Lite Model\n",
"\n",
"Convert the trained model to TensorFlow Lite model format with [metadata](https://www.tensorflow.org/lite/models/convert/metadata) so that you can later use in an on-device ML application. The label file and the vocab file are embedded in metadata. The default TFLite filename is `model.tflite`.\n",
"\n",
"In many on-device ML application, the model size is an important factor. Therefore, it is recommended that you apply quantize the model to make it smaller and potentially run faster.\n",
"The default post-training quantization technique is full integer quantization for the image classification task."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Im6wA9lK3TQB"
},
"outputs": [],
"source": [
"model.export(export_dir='.')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ROS2Ay2jMPCl"
},
"source": [
"See the image classification [examples guide](https://www.tensorflow.org/lite/examples/image_classification/overview) for more details about how to integrate the TensorFlow Lite model into mobile apps.\n",
"\n",
"This model can be integrated into an Android or an iOS app using the [ImageClassifier API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/image_classifier) of the [TensorFlow Lite Task Library](https://www.tensorflow.org/lite/inference_with_metadata/task_library/overview)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "habFnvRxxQ4A"
},
"source": [
"The allowed export formats can be one or a list of the following:\n",
"\n",
"* `ExportFormat.TFLITE`\n",
"* `ExportFormat.LABEL`\n",
"* `ExportFormat.SAVED_MODEL`\n",
"\n",
"By default, it just exports TensorFlow Lite model with metadata. You can also selectively export different files. For instance, exporting only the label file as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "BvxWsOTmKG4P"
},
"outputs": [],
"source": [
"model.export(export_dir='.', export_format=ExportFormat.LABEL)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-4jQaxyT5_KV"
},
"source": [
"You can also evaluate the tflite model with the `evaluate_tflite` method."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "S1YoPX5wOK-u"
},
"outputs": [],
"source": [
"model.evaluate_tflite('model.tflite', test_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "zNDBP2qA54aK"
},
"source": [
"## Advanced Usage\n",
"\n",
"The `create` function is the critical part of this library. It uses transfer learning with a pretrained model similar to the [tutorial](https://www.tensorflow.org/tutorials/images/transfer_learning).\n",
"\n",
"The `create` function contains the following steps:\n",
"\n",
"1. Split the data into training, validation, testing data according to parameter `validation_ratio` and `test_ratio`. The default value of `validation_ratio` and `test_ratio` are `0.1` and `0.1`.\n",
"2. Download a [Image Feature Vector](https://www.tensorflow.org/hub/common_signatures/images#image_feature_vector) as the base model from TensorFlow Hub. The default pre-trained model is EfficientNet-Lite0.\n",
"3. Add a classifier head with a Dropout Layer with `dropout_rate` between head layer and pre-trained model. The default `dropout_rate` is the default `dropout_rate` value from [make_image_classifier_lib](https://github.com/tensorflow/hub/blob/master/tensorflow_hub/tools/make_image_classifier/make_image_classifier_lib.py#L55) by TensorFlow Hub.\n",
"4. Preprocess the raw input data. Currently, preprocessing steps including normalizing the value of each image pixel to model input scale and resizing it to model input size. EfficientNet-Lite0 have the input scale `[0, 1]` and the input image size `[224, 224, 3]`.\n",
"5. Feed the data into the classifier model. By default, the training parameters such as training epochs, batch size, learning rate, momentum are the default values from [make_image_classifier_lib](https://github.com/tensorflow/hub/blob/master/tensorflow_hub/tools/make_image_classifier/make_image_classifier_lib.py#L55) by TensorFlow Hub. Only the classifier head is trained.\n",
"\n",
"\n",
"In this section, we describe several advanced topics, including switching to a different image classification model, changing the training hyperparameters etc.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Gc4Jk8TvBQfm"
},
"source": [
"## Customize Post-training quantization on the TensorFLow Lite model\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tD8BOYrHBiDt"
},
"source": [
"[Post-training quantization](https://www.tensorflow.org/lite/performance/post_training_quantization) is a conversion technique that can reduce model size and inference latency, while also improving CPU and hardware accelerator inference speed, with a little degradation in model accuracy. Thus, it's widely used to optimize the model.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "iyIo0d5TCzE2"
},
"source": [
"Model Maker library applies a default post-training quantization techique when exporting the model. If you want to customize post-training quantization, Model Maker supports multiple post-training quantization options using [QuantizationConfig](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/config/QuantizationConfig) as well. Let's take float16 quantization as an instance. First, define the quantization config."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "k8hL2mstCxQl"
},
"outputs": [],
"source": [
"config = QuantizationConfig.for_float16()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "K1gzx_rmFMOA"
},
"source": [
"Then we export the TensorFlow Lite model with such configuration."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WTJzFQnJFMjr"
},
"outputs": [],
"source": [
"model.export(export_dir='.', tflite_filename='model_fp16.tflite', quantization_config=config)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Safo0e40wKZW"
},
"source": [
"In Colab, you can download the model named `model_fp16.tflite` from the left sidebar, same as the uploading part mentioned above."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "A4kiTJtZ_sDm"
},
"source": [
"## Change the model\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "794vgj6ud7Ep"
},
"source": [
"### Change to the model that's supported in this library.\n",
"\n",
"This library supports EfficientNet-Lite models, MobileNetV2, ResNet50 by now. [EfficientNet-Lite](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/lite) are a family of image classification models that could achieve state-of-art accuracy and suitable for Edge devices. The default model is EfficientNet-Lite0.\n",
"\n",
"We could switch model to MobileNetV2 by just setting parameter `model_spec` to the MobileNetV2 model specification in `create` method."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7JKsJ6-P6ae1"
},
"outputs": [],
"source": [
"model = image_classifier.create(train_data, model_spec=model_spec.get('mobilenet_v2'), validation_data=validation_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gm_B1Wv08AxR"
},
"source": [
"Evaluate the newly retrained MobileNetV2 model to see the accuracy and loss in testing data."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "lB2Go3HW8X7_"
},
"outputs": [],
"source": [
"loss, accuracy = model.evaluate(test_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vAciGzVWtmWp"
},
"source": [
"### Change to the model in TensorFlow Hub\n",
"\n",
"Moreover, we could also switch to other new models that inputs an image and outputs a feature vector with TensorFlow Hub format.\n",
"\n",
"As [Inception V3](https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1) model as an example, we could define `inception_v3_spec` which is an object of [image_classifier.ModelSpec](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/image_classifier/ModelSpec) and contains the specification of the Inception V3 model.\n",
"\n",
"We need to specify the model name `name`, the url of the TensorFlow Hub model `uri`. Meanwhile, the default value of `input_image_shape` is `[224, 224]`. We need to change it to `[299, 299]` for Inception V3 model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "xdiMF2WMfAR4"
},
"outputs": [],
"source": [
"inception_v3_spec = image_classifier.ModelSpec(\n",
" uri='https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1')\n",
"inception_v3_spec.input_image_shape = [299, 299]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "T_GGIoXZCs5F"
},
"source": [
"Then, by setting parameter `model_spec` to `inception_v3_spec` in `create` method, we could retrain the Inception V3 model.\n",
"\n",
"The remaining steps are exactly same and we could get a customized InceptionV3 TensorFlow Lite model in the end."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "UhZ5IRKdeex3"
},
"source": [
"### Change your own custom model"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "svTjlZhrCrcV"
},
"source": [
"If we'd like to use the custom model that's not in TensorFlow Hub, we should create and export [ModelSpec](https://www.tensorflow.org/hub/api_docs/python/hub/ModuleSpec) in TensorFlow Hub.\n",
"\n",
"Then start to define `ModelSpec` object like the process above."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "4M9bn703AHt2"
},
"source": [
"## Change the training hyperparameters\n",
"We could also change the training hyperparameters like `epochs`, `dropout_rate` and `batch_size` that could affect the model accuracy. The model parameters you can adjust are:\n",
"\n",
"\n",
"* `epochs`: more epochs could achieve better accuracy until it converges but training for too many epochs may lead to overfitting.\n",
"* `dropout_rate`: The rate for dropout, avoid overfitting. None by default.\n",
"* `batch_size`: number of samples to use in one training step. None by default.\n",
"* `validation_data`: Validation data. If None, skips validation process. None by default.\n",
"* `train_whole_model`: If true, the Hub module is trained together with the classification layer on top. Otherwise, only train the top classification layer. None by default.\n",
"* `learning_rate`: Base learning rate. None by default.\n",
"* `momentum`: a Python float forwarded to the optimizer. Only used when\n",
" `use_hub_library` is True. None by default.\n",
"* `shuffle`: Boolean, whether the data should be shuffled. False by default.\n",
"* `use_augmentation`: Boolean, use data augmentation for preprocessing. False by default.\n",
"* `use_hub_library`: Boolean, use `make_image_classifier_lib` from tensorflow hub to retrain the model. This training pipeline could achieve better performance for complicated dataset with many categories. True by default. \n",
"* `warmup_steps`: Number of warmup steps for warmup schedule on learning rate. If None, the default warmup_steps is used which is the total training steps in two epochs. Only used when `use_hub_library` is False. None by default.\n",
"* `model_dir`: Optional, the location of the model checkpoint files. Only used when `use_hub_library` is False. None by default.\n",
"\n",
"Parameters which are None by default like `epochs` will get the concrete default parameters in [make_image_classifier_lib](https://github.com/tensorflow/hub/blob/02ab9b7d3455e99e97abecf43c5d598a5528e20c/tensorflow_hub/tools/make_image_classifier/make_image_classifier_lib.py#L54) from TensorFlow Hub library or [train_image_classifier_lib](https://github.com/tensorflow/examples/blob/f0260433d133fd3cea4a920d1e53ecda07163aee/tensorflow_examples/lite/model_maker/core/task/train_image_classifier_lib.py#L61).\n",
"\n",
"For example, we could train with more epochs.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "A3k7mhH54QcK"
},
"outputs": [],
"source": [
"model = image_classifier.create(train_data, validation_data=validation_data, epochs=10)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VaYBQymQDsXU"
},
"source": [
"Evaluate the newly retrained model with 10 training epochs."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "VafIYpKWD4Sw"
},
"outputs": [],
"source": [
"loss, accuracy = model.evaluate(test_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dhBU5NCy5Ji2"
},
"source": [
"# Read more\n",
"\n",
"You can read our [image classification](https://www.tensorflow.org/lite/examples/image_classification/overview) example to learn technical details. For more information, please refer to:\n",
"\n",
"* TensorFlow Lite Model Maker [guide](https://www.tensorflow.org/lite/models/modify/model_maker) and [API reference](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker).\n",
"* Task Library: [ImageClassifier](https://www.tensorflow.org/lite/inference_with_metadata/task_library/image_classifier) for deployment.\n",
"* The end-to-end reference apps: [Android](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/android), [iOS](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/ios), and [Raspberry PI](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/raspberry_pi).\n",
"\n"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "Model Maker Image Classification Tutorial",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
@@ -0,0 +1,91 @@
# TensorFlow Lite Model Maker
## Overview
The TensorFlow Lite Model Maker library simplifies the process of training a
TensorFlow Lite model using custom dataset. It uses transfer learning to reduce
the amount of training data required and shorten the training time.
## Supported Tasks
The Model Maker library currently supports the following ML tasks. Click the
links below for guides on how to train the model.
Supported Tasks | Task Utility
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------
Image Classification: [tutorial](https://www.tensorflow.org/lite/models/modify/model_maker/image_classification), [api](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/image_classifier) | Classify images into predefined categories.
Object Detection: [tutorial](https://www.tensorflow.org/lite/models/modify/model_maker/object_detection), [api](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/object_detector) | Detect objects in real time.
Text Classification: [tutorial](https://ai.google.dev/edge/litert/libraries/modify/text_classification), [api](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/text_classifier) | Classify text into predefined categories.
BERT Question Answer: [tutorial](https://www.tensorflow.org/lite/models/modify/model_maker/question_answer), [api](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/question_answer) | Find the answer in a certain context for a given question with BERT.
Audio Classification: [tutorial](https://www.tensorflow.org/lite/models/modify/model_maker/audio_classification), [api](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/audio_classifier) | Classify audio into predefined categories.
Recommendation: [demo](https://github.com/tensorflow/examples/blob/master/tensorflow_examples/lite/model_maker/demo/recommendation_demo.py), [api](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/recommendation) | Recommend items based on the context information for on-device scenario.
Searcher: [tutorial](https://www.tensorflow.org/lite/models/modify/model_maker/text_searcher), [api](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/searcher) | Search for similar text or image in a database.
If your tasks are not supported, please first use [TensorFlow](https://www.tensorflow.org/guide)
to retrain a TensorFlow model with transfer learning (following guides like
[images](https://www.tensorflow.org/tutorials/images/transfer_learning),
[text](https://www.tensorflow.org/official_models/fine_tuning_bert),
[audio](https://www.tensorflow.org/tutorials/audio/transfer_learning_audio)) or
train it from scratch, and then [convert](https://www.tensorflow.org/lite/convert)
it to TensorFlow Lite model.
## End-to-End Example
Model Maker allows you to train a TensorFlow Lite model using custom datasets in
just a few lines of code. For example, here are the steps to train an image
classification model.
```python
from tflite_model_maker import image_classifier
from tflite_model_maker.image_classifier import DataLoader
# Load input data specific to an on-device ML app.
data = DataLoader.from_folder('flower_photos/')
train_data, test_data = data.split(0.9)
# Customize the TensorFlow model.
model = image_classifier.create(train_data)
# Evaluate the model.
loss, accuracy = model.evaluate(test_data)
# Export to Tensorflow Lite model and label file in `export_dir`.
model.export(export_dir='/tmp/')
```
For more details, see the
[image classification guide](https://www.tensorflow.org/lite/models/modify/model_maker/image_classification).
## Installation
There are two ways to install Model Maker.
* Install a prebuilt pip package.
```shell
pip install tflite-model-maker
```
If you want to install nightly version, please follow the command:
```shell
pip install tflite-model-maker-nightly
```
* Clone the source code from GitHub and install.
```shell
git clone https://github.com/tensorflow/examples
cd examples/tensorflow_examples/lite/model_maker/pip_package
pip install -e .
```
TensorFlow Lite Model Maker depends on TensorFlow
[pip package](https://www.tensorflow.org/install/pip). For GPU drivers, please
refer to TensorFlow's [GPU guide](https://www.tensorflow.org/install/gpu) or
[installation guide](https://www.tensorflow.org/install).
## Python API Reference
You can find out Model Maker's public APIs in
[API reference](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker).
@@ -0,0 +1,863 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "h2q27gKz1H20"
},
"source": [
"##### Copyright 2024 The AI Edge Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "TUfAcER1oUS6"
},
"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": "Gb7qyhNL1yWt"
},
"source": [
"# Object Detection with TensorFlow Lite Model Maker"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Fw5Y7snSuG51"
},
"source": [
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://www.tensorflow.org/lite/models/modify/model_maker/object_detection\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/object_detection.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/object_detection.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/models/modify/model_maker/object_detection.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
" </td>\n",
"</table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "sr3q-gvm3cI8"
},
"source": [
"In this colab notebook, you'll learn how to use the [TensorFlow Lite Model Maker](https://www.tensorflow.org/lite/models/modify/model_maker) library to train a custom object detection model capable of detecting salads within images on a mobile device.\n",
"\n",
"The Model Maker library uses *transfer learning* to simplify the process of training a TensorFlow Lite model using a custom dataset. Retraining a TensorFlow Lite model with your own custom dataset reduces the amount of training data required and will shorten the training time.\n",
"\n",
"You'll use the publicly available *Salads* dataset, which was created from the [Open Images Dataset V4](https://storage.googleapis.com/openimages/web/index.html).\n",
"\n",
"Each image in the dataset \bcontains objects labeled as one of the following classes:\n",
"* Baked Good\n",
"* Cheese\n",
"* Salad\n",
"* Seafood\n",
"* Tomato\n",
"\n",
"The dataset contains the bounding-boxes specifying where each object locates, together with the object's label.\n",
"\n",
"Here is an example image from the dataset:\n",
"\n",
"<br/>\n",
"\n",
"<img src=\"https://cloud.google.com/vision/automl/object-detection/docs/images/quickstart-preparing_a_dataset.png\" width=\"400\" hspace=\"0\">\n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bcLF2PKkSbV3"
},
"source": [
"## Prerequisites\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2vvAObmTqglq"
},
"source": [
"### Install the required packages\n",
"Start by installing the required packages, including the Model Maker package from the [GitHub repo](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker) and the pycocotools library you'll use for evaluation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "qhl8lqVamEty"
},
"outputs": [],
"source": [
"!sudo apt -y install libportaudio2\n",
"!pip install -q --use-deprecated=legacy-resolver tflite-model-maker\n",
"!pip install -q pycocotools\n",
"!pip install -q opencv-python-headless==4.1.2.30\n",
"!pip uninstall -y tensorflow && pip install -q tensorflow==2.8.0"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "l6lRhVK9Q_0U"
},
"source": [
"Import the required packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "XtxiUeZEiXpt"
},
"outputs": [],
"source": [
"import numpy as np\n",
"import os\n",
"\n",
"from tflite_model_maker.config import QuantizationConfig\n",
"from tflite_model_maker.config import ExportFormat\n",
"from tflite_model_maker import model_spec\n",
"from tflite_model_maker import object_detector\n",
"\n",
"import tensorflow as tf\n",
"assert tf.__version__.startswith('2')\n",
"\n",
"tf.get_logger().setLevel('ERROR')\n",
"from absl import logging\n",
"logging.set_verbosity(logging.ERROR)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BRd13bfetO7B"
},
"source": [
"### Prepare the dataset\n",
"\n",
"Here you'll use the same dataset as the AutoML [quickstart](https://cloud.google.com/vision/automl/object-detection/docs/edge-quickstart#preparing_a_dataset).\n",
"\n",
"The *Salads* dataset is available at:\n",
" `gs://cloud-ml-data/img/openimage/csv/salads_ml_use.csv`.\n",
"\n",
"It contains 175 images for training, 25 images for validation, and 25 images for testing. The dataset has five classes: `Salad`, `Seafood`, `Tomato`, `Baked goods`, `Cheese`.\n",
"\n",
"<br/>\n",
"\n",
"The dataset is provided in CSV format:\n",
"```\n",
"TRAINING,gs://cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg,Salad,0.0,0.0954,,,0.977,0.957,,\n",
"VALIDATION,gs://cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg,Seafood,0.0154,0.1538,,,1.0,0.802,,\n",
"TEST,gs://cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg,Tomato,0.0,0.655,,,0.231,0.839,,\n",
"```\n",
"\n",
"* Each row corresponds to an object localized inside a larger image, with each object specifically designated as test, train, or validation data. You'll learn more about what that means in a later stage in this notebook.\n",
"* The three lines included here indicate **three distinct objects located inside the same image** available at `gs://cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg`.\n",
"* Each row has a different label: `Salad`, `Seafood`, `Tomato`, etc.\n",
"* Bounding boxes are specified for each image using the top left and bottom right vertices.\n",
"\n",
"Here is a visualization of these three lines:\n",
"\n",
"<br>\n",
"\n",
"<img src=\"https://cloud.google.com/vision/automl/object-detection/docs/images/quickstart-preparing_a_dataset.png\" width=\"400\" hspace=\"100\">\n",
"\n",
"If you want to know more about how to prepare your own CSV file and the minimum requirements for creating a valid dataset, see the [Preparing your training data](https://cloud.google.com/vision/automl/object-detection/docs/prepare) guide for more details.\n",
"\n",
"If you are new to Google Cloud, you may wonder what the `gs://` URL means. They are URLs of files stored on [Google Cloud Storage](https://cloud.google.com/storage) (GCS). If you make your files on GCS public or [authenticate your client](https://cloud.google.com/storage/docs/authentication#libauth), Model Maker can read those files similarly to your local files.\n",
"\n",
"However, you don't need to keep your images on Google Cloud to use Model Maker. You can use a local path in your CSV file and Model Maker will just work."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xushUyZXqP59"
},
"source": [
"## Quickstart"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vn61LJ9QbOPi"
},
"source": [
"There are six steps to training an object detection model:\n",
"\n",
"**Step 1. Choose an object detection model architecture.**\n",
"\n",
"This tutorial uses the EfficientDet-Lite0 model. EfficientDet-Lite[0-4] are a family of mobile/IoT-friendly object detection models derived from the [EfficientDet](https://arxiv.org/abs/1911.09070) architecture.\n",
"\n",
"Here is the performance of each EfficientDet-Lite models compared to each others.\n",
"\n",
"| Model architecture | Size(MB)* | Latency(ms)** | Average Precision*** |\n",
"|--------------------|-----------|---------------|----------------------|\n",
"| EfficientDet-Lite0 | 4.4 | 37 | 25.69% |\n",
"| EfficientDet-Lite1 | 5.8 | 49 | 30.55% |\n",
"| EfficientDet-Lite2 | 7.2 | 69 | 33.97% |\n",
"| EfficientDet-Lite3 | 11.4 | 116 | 37.70% |\n",
"| EfficientDet-Lite4 | 19.9 | 260 | 41.96% |\n",
"\n",
"<i> * Size of the integer quantized models. <br/>\n",
"** Latency measured on Pixel 4 using 4 threads on CPU. <br/>\n",
"*** Average Precision is the mAP (mean Average Precision) on the COCO 2017 validation dataset.\n",
"</i>\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "CtdZ-JDwMimd"
},
"outputs": [],
"source": [
"spec = model_spec.get('efficientdet_lite0')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "s5U-A3tw6Y27"
},
"source": [
"**Step 2. Load the dataset.**\n",
"\n",
"Model Maker will take input data in the CSV format. Use the `object_detector.DataLoader.from_csv` method to load the dataset and split them into the training, validation and test images.\n",
"\n",
"* Training images: These images are used to train the object detection model to recognize salad ingredients.\n",
"* Validation images: These are images that the model didn't see during the training process. You'll use them to decide when you should stop the training, to avoid [overfitting](https://en.wikipedia.org/wiki/Overfitting).\n",
"* Test images: These images are used to evaluate the final model performance.\n",
"\n",
"You can load the CSV file directly from Google Cloud Storage, but you don't need to keep your images on Google Cloud to use Model Maker. You can specify a local CSV file on your computer, and Model Maker will work just fine."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "HD5BvzWe6YKa"
},
"outputs": [],
"source": [
"train_data, validation_data, test_data = object_detector.DataLoader.from_csv('gs://cloud-ml-data/img/openimage/csv/salads_ml_use.csv')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2uZkLR6N6gDR"
},
"source": [
"**Step 3. Train the TensorFlow model with the training data.**\n",
"\n",
"* The EfficientDet-Lite0 model uses `epochs = 50` by default, which means it will go through the training dataset 50 times. You can look at the validation accuracy during training and stop early to avoid overfitting.\n",
"* Set `batch_size = 8` here so you will see that it takes 21 steps to go through the 175 images in the training dataset.\n",
"* Set `train_whole_model=True` to fine-tune the whole model instead of just training the head layer to improve accuracy. The trade-off is that it may take longer to train the model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "kwlYdTcg63xy"
},
"outputs": [],
"source": [
"model = object_detector.create(train_data, model_spec=spec, batch_size=8, train_whole_model=True, validation_data=validation_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-BzCHLWJ6h7q"
},
"source": [
"**Step 4. Evaluate the model with the test data.**\n",
"\n",
"After training the object detection model using the images in the training dataset, use the remaining 25 images in the test dataset to evaluate how the model performs against new data it has never seen before.\n",
"\n",
"As the default batch size is 64, it will take 1 step to go through the 25 images in the test dataset.\n",
"\n",
"The evaluation metrics are same as [COCO](https://cocodataset.org/#detection-eval)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8xmnl6Yy7ARn"
},
"outputs": [],
"source": [
"model.evaluate(test_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "CgCDMe0e6jlT"
},
"source": [
"**Step 5. Export as a TensorFlow Lite model.**\n",
"\n",
"Export the trained object detection model to the TensorFlow Lite format by specifying which folder you want to export the quantized model to. The default post-training quantization technique is full integer quantization."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Hm_UULdW7A9T"
},
"outputs": [],
"source": [
"model.export(export_dir='.')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ZQpahAIBqBPp"
},
"source": [
"**Step 6. Evaluate the TensorFlow Lite model.**\n",
"\n",
"Several factors can affect the model accuracy when exporting to TFLite:\n",
"* [Quantization](https://www.tensorflow.org/lite/performance/model_optimization) helps shrinking the model size by 4 times at the expense of some accuracy drop.\n",
"* The original TensorFlow model uses per-class [non-max supression (NMS)](https://www.coursera.org/lecture/convolutional-neural-networks/non-max-suppression-dvrjH) for post-processing, while the TFLite model uses global NMS that's much faster but less accurate.\n",
"Keras outputs maximum 100 detections while tflite outputs maximum 25 detections.\n",
"\n",
"Therefore you'll have to evaluate the exported TFLite model and compare its accuracy with the original TensorFlow model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "RS3Ell_lqH4e"
},
"outputs": [],
"source": [
"model.evaluate_tflite('model.tflite', test_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rVxaf3x_7OfB"
},
"source": [
"You can download the TensorFlow Lite model file using the left sidebar of Colab. Right-click on the `model.tflite` file and choose `Download` to download it to your local computer.\n",
"\n",
"This model can be integrated into an Android or an iOS app using the [ObjectDetector API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/object_detector) of the [TensorFlow Lite Task Library](https://www.tensorflow.org/lite/inference_with_metadata/task_library/overview).\n",
"\n",
"See the [TFLite Object Detection sample app](https://github.com/tensorflow/examples/tree/master/lite/examples/object_detection/android) for more details on how the model is used in an working app.\n",
"\n",
"*Note: Android Studio Model Binding does not support object detection yet so please use the TensorFlow Lite Task Library.*"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "me6_RwPZqNhX"
},
"source": [
"## (Optional) Test the TFLite model on your image\n",
"\n",
"You can test the trained TFLite model using images from the internet.\n",
"* Replace the `INPUT_IMAGE_URL` below with your desired input image.\n",
"* Adjust the `DETECTION_THRESHOLD` to change the sensitivity of the model. A lower threshold means the model will pickup more objects but there will also be more false detection. Meanwhile, a higher threshold means the model will only pickup objects that it has confidently detected.\n",
"\n",
"Although it requires some of boilerplate code to run the model in Python at this moment, integrating the model into a mobile app only requires a few lines of code."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "XqS0rFCrqM1o"
},
"outputs": [],
"source": [
"#@title Load the trained TFLite model and define some visualization functions\n",
"\n",
"import cv2\n",
"\n",
"from PIL import Image\n",
"\n",
"model_path = 'model.tflite'\n",
"\n",
"# Load the labels into a list\n",
"classes = ['???'] * model.model_spec.config.num_classes\n",
"label_map = model.model_spec.config.label_map\n",
"for label_id, label_name in label_map.as_dict().items():\n",
" classes[label_id-1] = label_name\n",
"\n",
"# Define a list of colors for visualization\n",
"COLORS = np.random.randint(0, 255, size=(len(classes), 3), dtype=np.uint8)\n",
"\n",
"def preprocess_image(image_path, input_size):\n",
" \"\"\"Preprocess the input image to feed to the TFLite model\"\"\"\n",
" img = tf.io.read_file(image_path)\n",
" img = tf.io.decode_image(img, channels=3)\n",
" img = tf.image.convert_image_dtype(img, tf.uint8)\n",
" original_image = img\n",
" resized_img = tf.image.resize(img, input_size)\n",
" resized_img = resized_img[tf.newaxis, :]\n",
" resized_img = tf.cast(resized_img, dtype=tf.uint8)\n",
" return resized_img, original_image\n",
"\n",
"\n",
"def detect_objects(interpreter, image, threshold):\n",
" \"\"\"Returns a list of detection results, each a dictionary of object info.\"\"\"\n",
"\n",
" signature_fn = interpreter.get_signature_runner()\n",
"\n",
" # Feed the input image to the model\n",
" output = signature_fn(images=image)\n",
"\n",
" # Get all outputs from the model\n",
" count = int(np.squeeze(output['output_0']))\n",
" scores = np.squeeze(output['output_1'])\n",
" classes = np.squeeze(output['output_2'])\n",
" boxes = np.squeeze(output['output_3'])\n",
"\n",
" results = []\n",
" for i in range(count):\n",
" if scores[i] >= threshold:\n",
" result = {\n",
" 'bounding_box': boxes[i],\n",
" 'class_id': classes[i],\n",
" 'score': scores[i]\n",
" }\n",
" results.append(result)\n",
" return results\n",
"\n",
"\n",
"def run_odt_and_draw_results(image_path, interpreter, threshold=0.5):\n",
" \"\"\"Run object detection on the input image and draw the detection results\"\"\"\n",
" # Load the input shape required by the model\n",
" _, input_height, input_width, _ = interpreter.get_input_details()[0]['shape']\n",
"\n",
" # Load the input image and preprocess it\n",
" preprocessed_image, original_image = preprocess_image(\n",
" image_path,\n",
" (input_height, input_width)\n",
" )\n",
"\n",
" # Run object detection on the input image\n",
" results = detect_objects(interpreter, preprocessed_image, threshold=threshold)\n",
"\n",
" # Plot the detection results on the input image\n",
" original_image_np = original_image.numpy().astype(np.uint8)\n",
" for obj in results:\n",
" # Convert the object bounding box from relative coordinates to absolute\n",
" # coordinates based on the original image resolution\n",
" ymin, xmin, ymax, xmax = obj['bounding_box']\n",
" xmin = int(xmin * original_image_np.shape[1])\n",
" xmax = int(xmax * original_image_np.shape[1])\n",
" ymin = int(ymin * original_image_np.shape[0])\n",
" ymax = int(ymax * original_image_np.shape[0])\n",
"\n",
" # Find the class index of the current object\n",
" class_id = int(obj['class_id'])\n",
"\n",
" # Draw the bounding box and label on the image\n",
" color = [int(c) for c in COLORS[class_id]]\n",
" cv2.rectangle(original_image_np, (xmin, ymin), (xmax, ymax), color, 2)\n",
" # Make adjustments to make the label visible for all objects\n",
" y = ymin - 15 if ymin - 15 > 15 else ymin + 15\n",
" label = \"{}: {:.0f}%\".format(classes[class_id], obj['score'] * 100)\n",
" cv2.putText(original_image_np, label, (xmin, y),\n",
" cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)\n",
"\n",
" # Return the final image\n",
" original_uint8 = original_image_np.astype(np.uint8)\n",
" return original_uint8"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "GkXtipXKqXp4"
},
"outputs": [],
"source": [
"#@title Run object detection and show the detection results\n",
"\n",
"INPUT_IMAGE_URL = \"https://storage.googleapis.com/cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg\" #@param {type:\"string\"}\n",
"DETECTION_THRESHOLD = 0.3 #@param {type:\"number\"}\n",
"\n",
"TEMP_FILE = '/tmp/image.png'\n",
"\n",
"!wget -q -O $TEMP_FILE $INPUT_IMAGE_URL\n",
"im = Image.open(TEMP_FILE)\n",
"im.thumbnail((512, 512), Image.ANTIALIAS)\n",
"im.save(TEMP_FILE, 'PNG')\n",
"\n",
"# Load the TFLite model\n",
"interpreter = tf.lite.Interpreter(model_path=model_path)\n",
"interpreter.allocate_tensors()\n",
"\n",
"# Run inference and draw detection result on the local copy of the original file\n",
"detection_result_image = run_odt_and_draw_results(\n",
" TEMP_FILE,\n",
" interpreter,\n",
" threshold=DETECTION_THRESHOLD\n",
")\n",
"\n",
"# Show the detection result\n",
"Image.fromarray(detection_result_image)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "oxgWQyYOqZha"
},
"source": [
"## (Optional) Compile For the Edge TPU\n",
"\n",
"Now that you have a quantized EfficientDet Lite model, it is possible to compile and deploy to a [Coral EdgeTPU](https://coral.ai/).\n",
"\n",
"**Step 1. Install the EdgeTPU Compiler**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Oy3QIn_YqaRP"
},
"outputs": [],
"source": [
"! curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -\n",
"\n",
"! echo \"deb https://packages.cloud.google.com/apt coral-edgetpu-stable main\" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list\n",
"\n",
"! sudo apt-get update\n",
"\n",
"! sudo apt-get install edgetpu-compiler"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qRWewhqFqeL_"
},
"source": [
"**Step 2. Select number of Edge TPUs, Compile**\n",
"\n",
"The EdgeTPU has 8MB of SRAM for caching model parameters ([more info](https://coral.ai/docs/edgetpu/compiler/#parameter-data-caching)). This means that for models that are larger than 8MB, inference time will be increased in order to transfer over model parameters. One way to avoid this is [Model Pipelining](https://coral.ai/docs/edgetpu/pipeline/) - splitting the model into segments that can have a dedicated EdgeTPU. This can significantly improve latency.\n",
"\n",
"The below table can be used as a reference for the number of Edge TPUs to use - the larger models will not compile for a single TPU as the intermediate tensors can't fit in on-chip memory.\n",
"\n",
"| Model architecture | Minimum TPUs | Recommended TPUs\n",
"|--------------------|-------|-------|\n",
"| EfficientDet-Lite0 | 1 | 1 |\n",
"| EfficientDet-Lite1 | 1 | 1 |\n",
"| EfficientDet-Lite2 | 1 | 2 |\n",
"| EfficientDet-Lite3 | 2 | 2 |\n",
"| EfficientDet-Lite4 | 2 | 3 |"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "LZdonJGCqieU"
},
"outputs": [],
"source": [
"NUMBER_OF_TPUS = 1#@param {type:\"number\"}\n",
"\n",
"!edgetpu_compiler model.tflite --num_segments=$NUMBER_OF_TPUS"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-g6_KQXnqlTC"
},
"source": [
"**Step 3. Download, Run Model**\n",
"\n",
"With the model(s) compiled, they can now be run on EdgeTPU(s) for object detection. First, download the compiled TensorFlow Lite model file using the left sidebar of Colab. Right-click on the `model_edgetpu.tflite` file and choose `Download` to download it to your local computer."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VkQFz_qzqrrA"
},
"source": [
"Now you can run the model in your preferred manner. Examples of detection include:\n",
"* [pycoral detection](https://github.com/google-coral/pycoral/blob/master/examples/detect_image.py)\n",
"* [Basic TFLite detection](https://github.com/google-coral/tflite/tree/master/python/examples/detection)\n",
"* [Example Video Detection](https://github.com/google-coral/examples-camera)\n",
"* [libcoral C++ API](https://github.com/google-coral/libcoral)\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EoWiA_zX8rxE"
},
"source": [
"## Advanced Usage\n",
"\n",
"This section covers advanced usage topics like adjusting the model and the training hyperparameters."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "p79NHCx0xFqb"
},
"source": [
"### Load the dataset\n",
"\n",
"#### Load your own data\n",
"\n",
"You can upload your own dataset to work through this tutorial. Upload your dataset by using the left sidebar in Colab.\n",
"\n",
"<img src=\"https://storage.googleapis.com/download.tensorflow.org/models/tflite/screenshots/model_maker_object_detection.png\" alt=\"Upload File\" width=\"1000\" hspace=\"0\">\n",
"\n",
"If you prefer not to upload your dataset to the cloud, you can also locally run the library by following the [guide](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker).\n",
"\n",
"#### Load your data with a different data format\n",
"\n",
"The Model Maker library also supports the `object_detector.DataLoader.from_pascal_voc` method to load data with [PASCAL VOC](https://towardsdatascience.com/coco-data-format-for-object-detection-a4c5eaf518c5#:~:text=Pascal%20VOC%20is%20an%20XML,for%20training%2C%20testing%20and%20validation) format. [makesense.ai](https://www.makesense.ai/) and [LabelImg](https://github.com/tzutalin/labelImg) are the tools that can annotate the image and save annotations as XML files in PASCAL VOC data format:\n",
"```python\n",
"object_detector.DataLoader.from_pascal_voc(image_dir, annotations_dir, label_map={1: \"person\", 2: \"notperson\"})\n",
"```\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "E8VxPiOLy4Gv"
},
"source": [
"### Customize the EfficientDet model hyperparameters\n",
"\n",
"The model and training pipeline parameters you can adjust are:\n",
"\n",
"* `model_dir`: The location to save the model checkpoint files. If not set, a temporary directory will be used.\n",
"* `steps_per_execution`: Number of steps per training execution.\n",
"* `moving_average_decay`: Float. The decay to use for maintaining moving averages of the trained parameters.\n",
"* `var_freeze_expr`: The regular expression to map the prefix name of variables to be frozen which means remaining the same during training. More specific, use `re.match(var_freeze_expr, variable_name)` in the codebase to map the variables to be frozen.\n",
"* `tflite_max_detections`: integer, 25 by default. The max number of output detections in the TFLite model.\n",
"* `strategy`: A string specifying which distribution strategy to use. Accepted values are 'tpu', 'gpus', None. tpu' means to use TPUStrategy. 'gpus' mean to use MirroredStrategy for multi-gpus. If None, use TF default with OneDeviceStrategy.\n",
"* `tpu`: The Cloud TPU to use for training. This should be either the name used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 url.\n",
"* `use_xla`: Use XLA even if strategy is not tpu. If strategy is tpu, always use XLA, and this flag has no effect.\n",
"* `profile`: Enable profile mode.\n",
"* `debug`: Enable debug mode.\n",
"\n",
"Other parameters that can be adjusted is shown in [hparams_config.py](https://github.com/google/automl/blob/df451765d467c5ed78bbdfd632810bc1014b123e/efficientdet/hparams_config.py#L170).\n",
"\n",
"\n",
"For instance, you can set the `var_freeze_expr='efficientnet'` which freezes the variables with name prefix `efficientnet` (default is `'(efficientnet|fpn_cells|resample_p6)'`). This allows the model to freeze untrainable variables and keep their value the same through training.\n",
"\n",
"```python\n",
"spec = model_spec.get('efficientdet_lite0')\n",
"spec.config.var_freeze_expr = 'efficientnet'\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "4J2qre1fwXsi"
},
"source": [
"### Change the Model Architecture\n",
"\n",
"You can change the model architecture by changing the `model_spec`. For instance, change the `model_spec` to the EfficientDet-Lite4 model.\n",
"\n",
"```python\n",
"spec = model_spec.get('efficientdet_lite4')\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LvQuy7RSDir3"
},
"source": [
"### Tune the training hyperparameters\n",
"\n",
"The `create` function is the driver function that the Model Maker library uses to create models. The `model_spec` parameter defines the model specification. The `object_detector.EfficientDetSpec` class is currently supported. The `create` function comprises of the following steps:\n",
"\n",
"1. Creates the model for the object detection according to `model_spec`.\n",
"2. Trains the model. The default epochs and the default batch size are set by the `epochs` and `batch_size` variables in the `model_spec` object.\n",
"You can also tune the training hyperparameters like `epochs` and `batch_size` that affect the model accuracy. For instance,\n",
"\n",
"* `epochs`: Integer, 50 by default. More epochs could achieve better accuracy, but may lead to overfitting.\n",
"* `batch_size`: Integer, 64 by default. The number of samples to use in one training step.\n",
"* `train_whole_model`: Boolean, False by default. If true, train the whole model. Otherwise, only train the layers that do not match `var_freeze_expr`.\n",
"\n",
"For example, you can train with less epochs and only the head layer. You can increase the number of epochs for better results.\n",
"\n",
"```python\n",
"model = object_detector.create(train_data, model_spec=spec, epochs=10, validation_data=validation_data)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3vPyZInPxJBT"
},
"source": [
"### Export to different formats"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0xqNIcBM-4YR"
},
"source": [
"The export formats can be one or a list of the following:\n",
"\n",
"* `ExportFormat.TFLITE`\n",
"* `ExportFormat.LABEL`\n",
"* `ExportFormat.SAVED_MODEL`\n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "enhsZhW3ApcX"
},
"source": [
"By default, it exports only the TensorFlow Lite model file containing the model [metadata](https://www.tensorflow.org/lite/models/convert/metadata) so that you can later use in an on-device ML application. The label file is embedded in metadata.\n",
"\n",
"In many on-device ML application, the model size is an important factor. Therefore, it is recommended that you quantize the model to make it smaller and potentially run faster. As for EfficientDet-Lite models, full integer quantization is used to quantize the model by default. Please refer to [Post-training quantization](https://www.tensorflow.org/lite/performance/post_training_quantization) for more detail.\n",
"\n",
"```python\n",
"model.export(export_dir='.')\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RLGZs6InAnP5"
},
"source": [
"You can also choose to export other files related to the model for better examination. For instance, exporting both the saved model and the label file as follows:\n",
"```python\n",
"model.export(export_dir='.', export_format=[ExportFormat.SAVED_MODEL, ExportFormat.LABEL])\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "W5q_McchQ2C4"
},
"source": [
"### Customize Post-training quantization on the TensorFlow Lite model\n",
"\n",
"[Post-training quantization](https://www.tensorflow.org/lite/performance/post_training_quantization) is a conversion technique that can reduce model size and inference latency, while also improving CPU and hardware accelerator inference speed, with a little degradation in model accuracy. Thus, it's widely used to optimize the model.\n",
"\n",
"Model Maker library applies a default post-training quantization technique when exporting the model. If you want to customize post-training quantization, Model Maker supports multiple post-training quantization options using [QuantizationConfig](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/config/QuantizationConfig) as well. Let's take float16 quantization as an instance. First, define the quantization config.\n",
"\n",
"```python\n",
"config = QuantizationConfig.for_float16()\n",
"```\n",
"\n",
"\n",
"Then we export the TensorFlow Lite model with such configuration.\n",
"\n",
"```python\n",
"model.export(export_dir='.', tflite_filename='model_fp16.tflite', quantization_config=config)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HS4u77W5gnzQ"
},
"source": [
"# Read more\n",
"\n",
"You can read our [object detection](https://www.tensorflow.org/lite/examples/object_detection/overview) example to learn technical details. For more information, please refer to:\n",
"\n",
"* TensorFlow Lite Model Maker [guide](https://www.tensorflow.org/lite/models/modify/model_maker) and [API reference](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker).\n",
"* Task Library: [ObjectDetector](https://www.tensorflow.org/lite/inference_with_metadata/task_library/object_detector) for deployment.\n",
"* The end-to-end reference apps: [Android](https://github.com/tensorflow/examples/tree/master/lite/examples/object_detection/android), [iOS](https://github.com/tensorflow/examples/tree/master/lite/examples/object_detection/ios), and [Raspberry PI](https://github.com/tensorflow/examples/tree/master/lite/examples/object_detection/raspberry_pi).\n",
"\n"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "Model Maker Object Detection Tutorial",
"provenance": [
{
"file_id": "1dbRXQCjtm-jBFC32DJ6YCVXnXBOG3M5t",
"timestamp": 1613441434239
},
{
"file_id": "https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/text_classification.ipynb",
"timestamp": 1612303859066
}
],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
@@ -0,0 +1,632 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "h2q27gKz1H20"
},
"source": [
"##### Copyright 2024 The AI Edge Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "TUfAcER1oUS6"
},
"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": "Gb7qyhNL1yWt"
},
"source": [
"# BERT Question Answer with TensorFlow Lite Model Maker"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Fw5Y7snSuG51"
},
"source": [
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://www.tensorflow.org/lite/models/modify/model_maker/question_answer\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/question_answer.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/question_answer.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/models/modify/model_maker/question_answer.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
" </td>\n",
"</table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "sr3q-gvm3cI8"
},
"source": [
"The [TensorFlow Lite Model Maker library](https://www.tensorflow.org/lite/models/modify/model_maker) simplifies the process of adapting and converting a TensorFlow model to particular input data when deploying this model for on-device ML applications.\n",
"\n",
"This notebook shows an end-to-end example that utilizes the Model Maker library to illustrate the adaptation and conversion of a commonly-used question answer model for question answer task."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "UxEHFTk755qw"
},
"source": [
"# Introduction to BERT Question Answer Task"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cFbKTCF25-SG"
},
"source": [
"The supported task in this library is extractive question answer task, which means given a passage and a question, the answer is the span in the passage. The image below shows an example for question answer.\n",
"\n",
"\n",
"<p align=\"center\"><img src=\"https://storage.googleapis.com/download.tensorflow.org/models/tflite/screenshots/model_maker_squad_showcase.png\" width=\"500\"></p>\n",
"\n",
"<p align=\"center\">\n",
" <em>Answers are spans in the passage (image credit: <a href=\"https://rajpurkar.github.io/mlx/qa-and-squad/\">SQuAD blog</a>) </em>\n",
"</p>\n",
"\n",
"As for the model of question answer task, the inputs should be the passage and question pair that are already preprocessed, the outputs should be the start logits and end logits for each token in the passage.\n",
"The size of input could be set and adjusted according to the length of passage and question."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gb7P4WQta8Ub"
},
"source": [
"## End-to-End Overview\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "w7cIHjIfbDlG"
},
"source": [
"The following code snippet demonstrates how to get the model within a few lines of code. The overall process includes 5 steps: (1) choose a model, (2) load data, (3) retrain the model, (4) evaluate, and (5) export it to TensorFlow Lite format."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xQPdlxZBYuZG"
},
"source": [
"```python\n",
"# Chooses a model specification that represents the model.\n",
"spec = model_spec.get('mobilebert_qa')\n",
"\n",
"# Gets the training data and validation data.\n",
"train_data = DataLoader.from_squad(train_data_path, spec, is_training=True)\n",
"validation_data = DataLoader.from_squad(validation_data_path, spec, is_training=False)\n",
"\n",
"# Fine-tunes the model.\n",
"model = question_answer.create(train_data, model_spec=spec)\n",
"\n",
"# Gets the evaluation result.\n",
"metric = model.evaluate(validation_data)\n",
"\n",
"# Exports the model to the TensorFlow Lite format with metadata in the export directory.\n",
"model.export(export_dir)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "exScAdvBbNEi"
},
"source": [
"The following sections explain the code in more detail."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bcLF2PKkSbV3"
},
"source": [
"## Prerequisites\n",
"\n",
"To run this example, install the required packages, including the Model Maker package from the [GitHub repo](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "qhl8lqVamEty"
},
"outputs": [],
"source": [
"!sudo apt -y install libportaudio2\n",
"!pip install -q tflite-model-maker-nightly"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "l6lRhVK9Q_0U"
},
"source": [
"Import the required packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "XtxiUeZEiXpt"
},
"outputs": [],
"source": [
"import numpy as np\n",
"import os\n",
"\n",
"import tensorflow as tf\n",
"assert tf.__version__.startswith('2')\n",
"\n",
"from tflite_model_maker import model_spec\n",
"from tflite_model_maker import question_answer\n",
"from tflite_model_maker.config import ExportFormat\n",
"from tflite_model_maker.question_answer import DataLoader"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "l65ctmtW7_FF"
},
"source": [
"The \"End-to-End Overview\" demonstrates a simple end-to-end example. The following sections walk through the example step by step to show more detail."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kJ_B8fMDOhMR"
},
"source": [
"## Choose a model_spec that represents a model for question answer\n",
"\n",
"Each `model_spec` object represents a specific model for question answer. The Model Maker currently supports MobileBERT and BERT-Base models.\n",
"\n",
"Supported Model | Name of model_spec | Model Description\n",
"--- | --- | ---\n",
"[MobileBERT](https://arxiv.org/pdf/2004.02984.pdf) | 'mobilebert_qa' | 4.3x smaller and 5.5x faster than BERT-Base while achieving competitive results, suitable for on-device scenario.\n",
"[MobileBERT-SQuAD](https://arxiv.org/pdf/2004.02984.pdf) | 'mobilebert_qa_squad' | Same model architecture as MobileBERT model and the initial model is already retrained on [SQuAD1.1](https://rajpurkar.github.io/SQuAD-explorer/).\n",
"[BERT-Base](https://arxiv.org/pdf/1810.04805.pdf) | 'bert_qa' | Standard BERT model that widely used in NLP tasks.\n",
"\n",
"In this tutorial, [MobileBERT-SQuAD](https://arxiv.org/pdf/2004.02984.pdf) is used as an example. Since the model is already retrained on [SQuAD1.1](https://rajpurkar.github.io/SQuAD-explorer/), it could coverage faster for question answer task.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "vEAWuZQ1PFiX"
},
"outputs": [],
"source": [
"spec = model_spec.get('mobilebert_qa_squad')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ygEncJxtl-nQ"
},
"source": [
"## Load Input Data Specific to an On-device ML App and Preprocess the Data\n",
"\n",
"The [TriviaQA](https://nlp.cs.washington.edu/triviaqa/) is a reading comprehension dataset containing over 650K question-answer-evidence triples. In this tutorial, you will use a subset of this dataset to learn how to use the Model Maker library.\n",
"\n",
"To load the data, convert the TriviaQA dataset to the [SQuAD1.1](https://rajpurkar.github.io/SQuAD-explorer/) format by running the [converter Python script](https://github.com/mandarjoshi90/triviaqa#miscellaneous) with `--sample_size=8000` and a set of `web` data. Modify the conversion code a little bit by:\n",
"* Skipping the samples that couldn't find any answer in the context document;\n",
"* Getting the original answer in the context without uppercase or lowercase.\n",
"\n",
"Download the archived version of the already converted dataset."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7tOfUr2KlgpU"
},
"outputs": [],
"source": [
"train_data_path = tf.keras.utils.get_file(\n",
" fname='triviaqa-web-train-8000.json',\n",
" origin='https://storage.googleapis.com/download.tensorflow.org/models/tflite/dataset/triviaqa-web-train-8000.json')\n",
"validation_data_path = tf.keras.utils.get_file(\n",
" fname='triviaqa-verified-web-dev.json',\n",
" origin='https://storage.googleapis.com/download.tensorflow.org/models/tflite/dataset/triviaqa-verified-web-dev.json')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "UfZk8GNr_1nc"
},
"source": [
"You can also train the MobileBERT model with your own dataset. If you are running this notebook on Colab, upload your data by using the left sidebar.\n",
"\n",
"<img src=\"https://storage.googleapis.com/download.tensorflow.org/models/tflite/screenshots/model_maker_question_answer.png\" alt=\"Upload File\" width=\"800\" hspace=\"100\">\n",
"\n",
"If you prefer not to upload your data to the cloud, you can also run the library offline by following the [guide](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "E051HBUM5owi"
},
"source": [
"Use the `DataLoader.from_squad` method to load and preprocess the [SQuAD format](https://rajpurkar.github.io/SQuAD-explorer/) data according to a specific `model_spec`. You can use either SQuAD2.0 or SQuAD1.1 formats. Setting parameter `version_2_with_negative` as `True` means the formats is SQuAD2.0. Otherwise, the format is SQuAD1.1. By default, `version_2_with_negative` is `False`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "I_fOlZsklmlL"
},
"outputs": [],
"source": [
"train_data = DataLoader.from_squad(train_data_path, spec, is_training=True)\n",
"validation_data = DataLoader.from_squad(validation_data_path, spec, is_training=False)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AWuoensX4vDA"
},
"source": [
"## Customize the TensorFlow Model\n",
"\n",
"Create a custom question answer model based on the loaded data. The `create` function comprises the following steps:\n",
"\n",
"1. Creates the model for question answer according to `model_spec`.\n",
"2. Train the question answer model. The default epochs and the default batch size are set according to two variables `default_training_epochs` and `default_batch_size` in the `model_spec` object."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "TvYSUuJY3QxR"
},
"outputs": [],
"source": [
"model = question_answer.create(train_data, model_spec=spec)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0JKI-pNc8idH"
},
"source": [
"Have a look at the detailed model structure."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "gd7Hs8TF8n3H"
},
"outputs": [],
"source": [
"model.summary()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LP5FPk_tOxoZ"
},
"source": [
"## Evaluate the Customized Model\n",
"\n",
"Evaluate the model on the validation data and get a dict of metrics including `f1` score and `exact match` etc. Note that metrics are different for SQuAD1.1 and SQuAD2.0."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "A8c2ZQ0J3Riy"
},
"outputs": [],
"source": [
"model.evaluate(validation_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aeHoGAceO2xV"
},
"source": [
"## Export to TensorFlow Lite Model\n",
"\n",
"Convert the trained model to TensorFlow Lite model format with [metadata](https://www.tensorflow.org/lite/models/convert/metadata) so that you can later use in an on-device ML application. The vocab file are embedded in metadata. The default TFLite filename is `model.tflite`.\n",
"\n",
"In many on-device ML application, the model size is an important factor. Therefore, it is recommended that you apply quantize the model to make it smaller and potentially run faster.\n",
"The default post-training quantization technique is dynamic range quantization for the BERT and MobileBERT models."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Im6wA9lK3TQB"
},
"outputs": [],
"source": [
"model.export(export_dir='.')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "w12kvDdHJIGH"
},
"source": [
"You can use the TensorFlow Lite model file in the [bert_qa](https://github.com/tensorflow/examples/tree/master/lite/examples/bert_qa/android) reference app using [BertQuestionAnswerer API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/bert_question_answerer) in [TensorFlow Lite Task Library](https://www.tensorflow.org/lite/inference_with_metadata/task_library/overview) by downloading it from the left sidebar on Colab."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VFnJPvq3VGh3"
},
"source": [
"The allowed export formats can be one or a list of the following:\n",
"\n",
"* `ExportFormat.TFLITE`\n",
"* `ExportFormat.VOCAB`\n",
"* `ExportFormat.SAVED_MODEL`\n",
"\n",
"By default, it just exports TensorFlow Lite model with metadata. You can also selectively export different files. For instance, exporting only the vocab file as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ro2hz4kXVImY"
},
"outputs": [],
"source": [
"model.export(export_dir='.', export_format=ExportFormat.VOCAB)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HZKYthlVrTos"
},
"source": [
"You can also evaluate the tflite model with the `evaluate_tflite` method. This step is expected to take a long time."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ochbq95ZrVFX"
},
"outputs": [],
"source": [
"model.evaluate_tflite('model.tflite', validation_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EoWiA_zX8rxE"
},
"source": [
"## Advanced Usage\n",
"\n",
"The `create` function is the critical part of this library in which the `model_spec` parameter defines the model specification. The `BertQASpec` class is currently supported. There are 2 models: MobileBERT model, BERT-Base model. The `create` function comprises the following steps:\n",
"\n",
"1. Creates the model for question answer according to `model_spec`.\n",
"2. Train the question answer model.\n",
"\n",
"This section describes several advanced topics, including adjusting the model, tuning the training hyperparameters etc."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "mwtiksguDfhl"
},
"source": [
"### Adjust the model\n",
"\n",
"You can adjust the model infrastructure like parameters `seq_len` and `query_len` in the `BertQASpec` class.\n",
"\n",
"Adjustable parameters for model:\n",
"\n",
"* `seq_len`: Length of the passage to feed into the model.\n",
"* `query_len`: Length of the question to feed into the model.\n",
"* `doc_stride`: The stride when doing a sliding window approach to take chunks of the documents.\n",
"* `initializer_range`: The stdev of the truncated_normal_initializer for initializing all weight matrices.\n",
"* `trainable`: Boolean, whether pre-trained layer is trainable.\n",
"\n",
"Adjustable parameters for training pipeline:\n",
"\n",
"* `model_dir`: The location of the model checkpoint files. If not set, temporary directory will be used.\n",
"* `dropout_rate`: The rate for dropout.\n",
"* `learning_rate`: The initial learning rate for Adam.\n",
"* `predict_batch_size`: Batch size for prediction.\n",
"* `tpu`: TPU address to connect to. Only used if using tpu.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cAOd5_bzH9AQ"
},
"source": [
"For example, you can train the model with a longer sequence length. If you change the model, you must first construct a new `model_spec`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "e9WBN0UTQoMN"
},
"outputs": [],
"source": [
"new_spec = model_spec.get('mobilebert_qa')\n",
"new_spec.seq_len = 512"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "6LSTdghTP0Cv"
},
"source": [
"The remaining steps are the same. Note that you must rerun both the `dataloader` and `create` parts as different model specs may have different preprocessing steps.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LvQuy7RSDir3"
},
"source": [
"### Tune training hyperparameters\n",
"You can also tune the training hyperparameters like `epochs` and `batch_size` to impact the model performance. For instance,\n",
"\n",
"* `epochs`: more epochs could achieve better performance, but may lead to overfitting.\n",
"* `batch_size`: number of samples to use in one training step.\n",
"\n",
"For example, you can train with more epochs and with a bigger batch size like:\n",
"\n",
"```python\n",
"model = question_answer.create(train_data, model_spec=spec, epochs=5, batch_size=64)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Eq6B9lKMfhS6"
},
"source": [
"### Change the Model Architecture\n",
"\n",
"You can change the base model your data trains on by changing the `model_spec`. For example, to change to the BERT-Base model, run:\n",
"\n",
"```python\n",
"spec = model_spec.get('bert_qa')\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "L2d7yycrgu6L"
},
"source": [
"The remaining steps are the same."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "wFQrDMXzOVoB"
},
"source": [
"### Customize Post-training quantization on the TensorFlow Lite model\n",
"\n",
"[Post-training quantization](https://www.tensorflow.org/lite/performance/post_training_quantization) is a conversion technique that can reduce model size and inference latency, while also improving CPU and hardware accelerator inference speed, with a little degradation in model accuracy. Thus, it's widely used to optimize the model.\n",
"\n",
"Model Maker library applies a default post-training quantization techique when exporting the model. If you want to customize post-training quantization, Model Maker supports multiple post-training quantization options using [QuantizationConfig](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/config/QuantizationConfig) as well. Let's take float16 quantization as an instance. First, define the quantization config.\n",
"\n",
"```python\n",
"config = QuantizationConfig.for_float16()\n",
"```\n",
"\n",
"\n",
"Then we export the TensorFlow Lite model with such configuration.\n",
"\n",
"```python\n",
"model.export(export_dir='.', tflite_filename='model_fp16.tflite', quantization_config=config)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "wPVopCeB6LV6"
},
"source": [
"# Read more\n",
"\n",
"You can read our [BERT Question and Answer](https://www.tensorflow.org/lite/examples/bert_qa/overview) example to learn technical details. For more information, please refer to:\n",
"\n",
"* TensorFlow Lite Model Maker [guide](https://www.tensorflow.org/lite/models/modify/model_maker) and [API reference](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker).\n",
"* Task Library: [BertQuestionAnswerer](https://www.tensorflow.org/lite/inference_with_metadata/task_library/bert_question_answerer) for deployment.\n",
"* The end-to-end reference apps: [Android](https://github.com/tensorflow/examples/tree/master/lite/examples/bert_qa/android) and [iOS](https://github.com/tensorflow/examples/tree/master/lite/examples/bert_qa/ios)."
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "Model Maker Question Answer Tutorial",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
@@ -0,0 +1,910 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "3XX46cTrh6iD"
},
"source": [
"##### Copyright 2024 The AI Edge Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "sKrlWr6Kh-mF"
},
"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": "hST65kOHXpiL"
},
"source": [
"# Retrain a speech recognition model with TensorFlow Lite Model Maker\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "nShlCXGkbRVA"
},
"source": [
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://www.tensorflow.org/lite/models/modify/model_maker/speech_recognition\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/speech_recognition.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/speech_recognition.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/models/modify/model_maker/speech_recognition.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
" </td>\n",
"\n",
"</table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BB5k6xNKJ5Xe"
},
"source": [
"\n",
"In this colab notebook, you'll learn how to use the [TensorFlow Lite Model Maker](https://www.tensorflow.org/lite/models/modify/model_maker) to train a speech recognition model that can classify spoken words or short phrases using one-second sound samples. The Model Maker library uses transfer learning to retrain an existing TensorFlow model with a new dataset, which reduces the amount of sample data and time required for training. \n",
"\n",
"By default, this notebook retrains the model (BrowserFft, from the [TFJS Speech Command Recognizer](https://github.com/tensorflow/tfjs-models/tree/master/speech-commands#speech-command-recognizer)) using a subset of words from the [speech commands dataset](https://www.tensorflow.org/datasets/catalog/speech_commands) (such as \"up,\" \"down,\" \"left,\" and \"right\"). Then it exports a TFLite model that you can run on a mobile device or embedded system (such as a Raspberry Pi). It also exports the trained model as a TensorFlow SavedModel.\n",
"\n",
"This notebook is also designed to accept a custom dataset of WAV files, uploaded to Colab in a ZIP file. The more samples you have for each class, the better your accuracy will be, but because the transfer learning process uses feature embeddings from the pre-trained model, you can still get a fairly accurate model with only a few dozen samples in each of your classes.\n",
"\n",
"**Note:** The model we'll be training is optimized for speech recognition with one-second samples. If you want to perform more generic audio classification (such as detecting different types of music), we suggest you instead follow [this Colab to retrain an audio classifier](https://colab.sandbox.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/audio_classification.ipynb).\n",
"\n",
"If you want to run the notebook with the default speech dataset, you can run the whole thing now by clicking **Runtime > Run all** in the Colab toolbar. However, if you want to use your own dataset, then continue down to [Prepare the dataset](#scrollTo=cBsSAeYLkc1Z) and follow the instructions there.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AeZZ_cSsZfPx"
},
"source": [
"### Import the required packages\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MelHQlE7FVue"
},
"source": [
"You'll need TensorFlow, TFLite Model Maker, and some modules for audio manipulation, playback, and visualizations."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "wbMc4vHjaYdQ"
},
"outputs": [],
"source": [
"!sudo apt -y install libportaudio2\n",
"!pip install tflite-model-maker"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rwUA9u4oWoCR"
},
"outputs": [],
"source": [
"import os\n",
"import glob\n",
"import random\n",
"import shutil\n",
"\n",
"import librosa\n",
"import soundfile as sf\n",
"from IPython.display import Audio\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"import seaborn as sns\n",
"\n",
"import tensorflow as tf\n",
"import tflite_model_maker as mm\n",
"from tflite_model_maker import audio_classifier\n",
"from tflite_model_maker.config import ExportFormat\n",
"\n",
"print(f\"TensorFlow Version: {tf.__version__}\")\n",
"print(f\"Model Maker Version: {mm.__version__}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cBsSAeYLkc1Z"
},
"source": [
"## Prepare the dataset"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "NTTSXJxJq2Bz"
},
"source": [
"To train with the default speech dataset, just run all the code below as-is.\n",
"\n",
"But if you want to train with your own speech dataset, follow these steps:\n",
"\n",
"**Note:** \n",
"The model you'll retrain expects input data to be roughly one second of audio at 44.1 kHz. Model Maker perfoms automatic resampling for the training dataset, so there's no need to resample your dataset if it has a sample rate other than 44.1 kHz. But beware that audio samples longer than one second will be split into multiple one-second chunks, and the final chunk will be discarded if it's shorter than one second.\n",
"\n",
"1. Be sure each sample in your dataset is in **WAV file format, about one second long**. Then create a ZIP file with all your WAV files, organized into separate subfolders for each classification. For example, each sample for a speech command \"yes\" should be in a subfolder named \"yes\". Even if you have only one class, the samples must be saved in a subdirectory with the class name as the directory name. (This script assumes your dataset **is not split** into train/validation/test sets and performs that split for you.)\n",
"2. Click the **Files** tab in the left panel and just drag-drop your ZIP file there to upload it.\n",
"3. Use the following drop-down option to set **`use_custom_dataset`** to True.\n",
"4. Then skip to [Prepare a custom audio dataset](#scrollTo=EobYerLQkiF1) to specify your ZIP filename and dataset directory name.\n",
"\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "AK9o98X7qyhU"
},
"outputs": [],
"source": [
"use_custom_dataset = False #@param [\"False\", \"True\"] {type:\"raw\"}"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "D2sNXbYVHjjy"
},
"source": [
"### Generate a background noise dataset"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gBVClNMwHtMD"
},
"source": [
"Whether you're using the default speech dataset or a custom dataset, you should have a good set of background noises so your model can distinguish speech from other noises (including silence). \n",
"\n",
"Because the following background samples are provided in WAV files that are a minute long or longer, we need to split them up into smaller one-second samples so we can reserve some for our test dataset. We'll also combine a couple different sample sources to build a comprehensive set of background noises and silence:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "qvJd9VfmHu29"
},
"outputs": [],
"source": [
"tf.keras.utils.get_file('speech_commands_v0.01.tar.gz',\n",
" 'http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz',\n",
" cache_dir='./',\n",
" cache_subdir='dataset-speech',\n",
" extract=True)\n",
"tf.keras.utils.get_file('background_audio.zip',\n",
" 'https://storage.googleapis.com/download.tensorflow.org/models/tflite/sound_classification/background_audio.zip',\n",
" cache_dir='./',\n",
" cache_subdir='dataset-background',\n",
" extract=True)\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7CAVFc3woB3_"
},
"source": [
"**Note:** Although there is a newer version available, we're using v0.01 of the speech commands dataset because it's a smaller download. v0.01 includes 30 commands, while v0.02 adds five more (\"backward\", \"forward\", \"follow\", \"learn\", and \"visual\")."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "xgwWNifGL-3b"
},
"outputs": [],
"source": [
"# Create a list of all the background wav files\n",
"files = glob.glob(os.path.join('./dataset-speech/_background_noise_', '*.wav'))\n",
"files = files + glob.glob(os.path.join('./dataset-background', '*.wav'))\n",
"\n",
"background_dir = './background'\n",
"os.makedirs(background_dir, exist_ok=True)\n",
"\n",
"# Loop through all files and split each into several one-second wav files\n",
"for file in files:\n",
" filename = os.path.basename(os.path.normpath(file))\n",
" print('Splitting', filename)\n",
" name = os.path.splitext(filename)[0]\n",
" rate = librosa.get_samplerate(file)\n",
" length = round(librosa.get_duration(filename=file))\n",
" for i in range(length - 1):\n",
" start = i * rate\n",
" stop = (i * rate) + rate\n",
" data, _ = sf.read(file, start=start, stop=stop)\n",
" sf.write(os.path.join(background_dir, name + str(i) + '.wav'), data, rate)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bVlvVq-SkeeO"
},
"source": [
"### Prepare the speech commands dataset"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "q_q22T5UHbJG"
},
"source": [
"We already downloaded the speech commands dataset, so now we just need to prune the number of classes for our model.\n",
"\n",
"This dataset includes over 30 speech command classifications, and most of them have over 2,000 samples. But because we're using transfer learning, we don't need that many samples. So the following code does a few things:\n",
"\n",
"+ Specify which classifications we want to use, and delete the rest.\n",
"+ Keep only 150 samples of each class for training (to prove that transfer learning works well with smaller datasets and simply to reduce the training time).\n",
"+ Create a separate directory for a test dataset so we can easily run inference with them later."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "zUSRpw2nOp8p"
},
"outputs": [],
"source": [
"if not use_custom_dataset:\n",
" commands = [ \"up\", \"down\", \"left\", \"right\", \"go\", \"stop\", \"on\", \"off\", \"background\"]\n",
" dataset_dir = './dataset-speech'\n",
" test_dir = './dataset-test'\n",
"\n",
" # Move the processed background samples\n",
" shutil.move(background_dir, os.path.join(dataset_dir, 'background')) \n",
"\n",
" # Delete all directories that are not in our commands list\n",
" dirs = glob.glob(os.path.join(dataset_dir, '*/'))\n",
" for dir in dirs:\n",
" name = os.path.basename(os.path.normpath(dir))\n",
" if name not in commands:\n",
" shutil.rmtree(dir)\n",
"\n",
" # Count is per class\n",
" sample_count = 150\n",
" test_data_ratio = 0.2\n",
" test_count = round(sample_count * test_data_ratio)\n",
"\n",
" # Loop through child directories (each class of wav files)\n",
" dirs = glob.glob(os.path.join(dataset_dir, '*/'))\n",
" for dir in dirs:\n",
" files = glob.glob(os.path.join(dir, '*.wav'))\n",
" random.seed(42)\n",
" random.shuffle(files)\n",
" # Move test samples:\n",
" for file in files[sample_count:sample_count + test_count]:\n",
" class_dir = os.path.basename(os.path.normpath(dir))\n",
" os.makedirs(os.path.join(test_dir, class_dir), exist_ok=True)\n",
" os.rename(file, os.path.join(test_dir, class_dir, os.path.basename(file)))\n",
" # Delete remaining samples\n",
" for file in files[sample_count + test_count:]:\n",
" os.remove(file)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EobYerLQkiF1"
},
"source": [
"### Prepare a custom dataset"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "f3xTvDP3knMd"
},
"source": [
"If you want to train the model with our own speech dataset, you need to upload your samples as WAV files in a ZIP ([as described above](#scrollTo=cBsSAeYLkc1Z)) and modify the following variables to specify your dataset:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "77PsQAKA4Arx"
},
"outputs": [],
"source": [
"if use_custom_dataset:\n",
" # Specify the ZIP file you uploaded:\n",
" !unzip YOUR-FILENAME.zip\n",
" # Specify the unzipped path to your custom dataset\n",
" # (this path contains all the subfolders with classification names):\n",
" dataset_dir = './YOUR-DIRNAME'"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hwp6EQqvttgf"
},
"source": [
"After changing the filename and path name above, you're ready to train the model with your custom dataset. In the Colab toolbar, select **Runtime > Run all** to run the whole notebook.\n",
"\n",
"The following code integrates our new background noise samples into your dataset and then separates a portion of all samples to create a test set."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tMQ6cpw_B9e_"
},
"outputs": [],
"source": [
"def move_background_dataset(dataset_dir):\n",
" dest_dir = os.path.join(dataset_dir, 'background')\n",
" if os.path.exists(dest_dir):\n",
" files = glob.glob(os.path.join(background_dir, '*.wav'))\n",
" for file in files:\n",
" shutil.move(file, dest_dir)\n",
" else:\n",
" shutil.move(background_dir, dest_dir)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "45iru8OdliG3"
},
"outputs": [],
"source": [
"if use_custom_dataset:\n",
" # Move background samples into custom dataset\n",
" move_background_dataset(dataset_dir)\n",
"\n",
" # Now we separate some of the files that we'll use for testing:\n",
" test_dir = './dataset-test'\n",
" test_data_ratio = 0.2\n",
" dirs = glob.glob(os.path.join(dataset_dir, '*/'))\n",
" for dir in dirs:\n",
" files = glob.glob(os.path.join(dir, '*.wav'))\n",
" test_count = round(len(files) * test_data_ratio)\n",
" random.seed(42)\n",
" random.shuffle(files)\n",
" # Move test samples:\n",
" for file in files[:test_count]:\n",
" class_dir = os.path.basename(os.path.normpath(dir))\n",
" os.makedirs(os.path.join(test_dir, class_dir), exist_ok=True)\n",
" os.rename(file, os.path.join(test_dir, class_dir, os.path.basename(file)))\n",
" print('Moved', test_count, 'images from', class_dir)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "myPa1dfEoagz"
},
"source": [
"### Play a sample"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "1jecBYREgMk6"
},
"source": [
"To be sure the dataset looks correct, let's play at a random sample from the test set:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "gLC3ayJsoeNw"
},
"outputs": [],
"source": [
"def get_random_audio_file(samples_dir):\n",
" files = os.path.abspath(os.path.join(samples_dir, '*/*.wav'))\n",
" files_list = glob.glob(files)\n",
" random_audio_path = random.choice(files_list)\n",
" return random_audio_path\n",
"\n",
"def show_sample(audio_path):\n",
" audio_data, sample_rate = sf.read(audio_path)\n",
" class_name = os.path.basename(os.path.dirname(audio_path))\n",
" print(f'Class: {class_name}')\n",
" print(f'File: {audio_path}')\n",
" print(f'Sample rate: {sample_rate}')\n",
" print(f'Sample length: {len(audio_data)}')\n",
"\n",
" plt.title(class_name)\n",
" plt.plot(audio_data)\n",
" display(Audio(audio_data, rate=sample_rate))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "todbtEWFy0mj"
},
"outputs": [],
"source": [
"random_audio = get_random_audio_file(test_dir)\n",
"show_sample(random_audio)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "f-jRIWcQv7xt"
},
"source": [
"## Define the model"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pQj1Mf7YZELS"
},
"source": [
"\n",
"When using Model Maker to retrain any model, you have to start by defining a model spec. The spec defines the base model from which your new model will extract feature embeddings to begin learning new classes. The spec for this speech recognizer is based on the pre-trained [BrowserFft model from TFJS](https://github.com/tensorflow/tfjs-models/tree/master/speech-commands#speech-command-recognizer).\n",
"\n",
"The model expects input as an audio sample that's 44.1 kHz, and just under a second long: the exact sample length must be 44034 frames.\n",
"\n",
"You don't need to do any resampling with your training dataset. Model Maker takes care of that for you. But when you later run inference, you must be sure that your input matches that expected format.\n",
"\n",
"All you need to do here is instantiate the [`BrowserFftSpec`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/audio_classifier/BrowserFftSpec):\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tUcxtfHXY7XS"
},
"outputs": [],
"source": [
"spec = audio_classifier.BrowserFftSpec()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "maTOoRvAwI9l"
},
"source": [
"## Load your dataset "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "UASCEHoVwQ1q"
},
"source": [
"Now you need to load your dataset according to the model specifications. Model Maker includes the [`DataLoader`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/audio_classifier/DataLoader) API, which will load your dataset from a folder and ensure it's in the expected format for the model spec.\n",
"\n",
"We already reserved some test files by moving them to a separate directory, which makes it easier to run inference with them later. Now we'll create a `DataLoader` for each split: the training set, the validation set, and the test set."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bAhAfHwiw2_F"
},
"source": [
"#### Load the speech commands dataset\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cX0RqETqZgzo"
},
"outputs": [],
"source": [
"if not use_custom_dataset:\n",
" train_data_ratio = 0.8\n",
" train_data = audio_classifier.DataLoader.from_folder(\n",
" spec, dataset_dir, cache=True)\n",
" train_data, validation_data = train_data.split(train_data_ratio)\n",
" test_data = audio_classifier.DataLoader.from_folder(\n",
" spec, test_dir, cache=True)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2OWQ_O9_t-C-"
},
"source": [
"#### Load a custom dataset"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "IPiwaJwMt7yo"
},
"source": [
"**Note:** Setting `cache=True` is important to make training faster (especially when the dataset must be re-sampled) but it will also require more RAM to hold the data. If you use a very large custom dataset, caching might exceed your RAM capacity."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "code",
"id": "e86Ej-ZmuCzy"
},
"outputs": [],
"source": [
"if use_custom_dataset:\n",
" train_data_ratio = 0.8\n",
" train_data = audio_classifier.DataLoader.from_folder(\n",
" spec, dataset_dir, cache=True)\n",
" train_data, validation_data = train_data.split(train_data_ratio)\n",
" test_data = audio_classifier.DataLoader.from_folder(\n",
" spec, test_dir, cache=True)\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hh1P_zfzwbfE"
},
"source": [
"## Train the model\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ziMghju-Rts2"
},
"source": [
"Now we'll use the Model Maker [`create()`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/audio_classifier/create) function to create a model based on our model spec and training dataset, and begin training.\n",
"\n",
"If you're using a custom dataset, you might want to change the batch size as appropriate for the number of samples in your train set.\n",
"\n",
"**Note:** The first epoch takes longer because it must create the cache. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "GYaZvaOPgLUC"
},
"outputs": [],
"source": [
"# If your dataset has fewer than 100 samples per class,\n",
"# you might want to try a smaller batch size\n",
"batch_size = 25\n",
"epochs = 25\n",
"model = audio_classifier.create(train_data, spec, validation_data, batch_size, epochs)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "mtLuRA2xweZA"
},
"source": [
"## Review the model performance"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "oXMEHZkAxJTl"
},
"source": [
"Even if the accuracy/loss looks good from the training output above, it's important to also run the model using test data that the model has not seen yet, which is what the `evaluate()` method does here:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "n_4MGpzhWVhr"
},
"outputs": [],
"source": [
"model.evaluate(test_data)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HqvpAnqsVExO"
},
"source": [
"### View the confusion matrix"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8QRRAM39aOxS"
},
"source": [
"When training a classification model such as this one, it's also useful to inspect the [confusion matrix](https://en.wikipedia.org/wiki/Confusion_matrix). The confusion matrix gives you detailed visual representation of how well your classifier performs for each classification in your test data."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "zqB3c0368iH3"
},
"outputs": [],
"source": [
"def show_confusion_matrix(confusion, test_labels):\n",
" \"\"\"Compute confusion matrix and normalize.\"\"\"\n",
" confusion_normalized = confusion.astype(\"float\") / confusion.sum(axis=1)\n",
" sns.set(rc = {'figure.figsize':(6,6)})\n",
" sns.heatmap(\n",
" confusion_normalized, xticklabels=test_labels, yticklabels=test_labels,\n",
" cmap='Blues', annot=True, fmt='.2f', square=True, cbar=False)\n",
" plt.title(\"Confusion matrix\")\n",
" plt.ylabel(\"True label\")\n",
" plt.xlabel(\"Predicted label\")\n",
"\n",
"confusion_matrix = model.confusion_matrix(test_data)\n",
"show_confusion_matrix(confusion_matrix.numpy(), test_data.index_to_label)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yASrikBgZ9ZO"
},
"source": [
"## Export the model\n",
"\n",
"The last step is exporting your model into the TensorFlow Lite format for execution on mobile/embedded devices and into the [SavedModel format](https://www.tensorflow.org/guide/saved_model) for execution elsewhere.\n",
"\n",
"When exporting a `.tflite` file from Model Maker, it includes [model metadata](https://www.tensorflow.org/lite/inference_with_metadata/overview) that describes various details that can later help during inference. It even includes a copy of the classification labels file, so you don't need to a separate `labels.txt` file. (In the next section, we show how to use this metadata to run an inference.)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4gEf59NfGWjq"
},
"outputs": [],
"source": [
"TFLITE_FILENAME = 'browserfft-speech.tflite'\n",
"SAVE_PATH = './models'"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Xw_ehPxAdQlz"
},
"outputs": [],
"source": [
"print(f'Exporing the model to {SAVE_PATH}')\n",
"model.export(SAVE_PATH, tflite_filename=TFLITE_FILENAME)\n",
"model.export(SAVE_PATH, export_format=[mm.ExportFormat.SAVED_MODEL, mm.ExportFormat.LABEL])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "lIC1ddGq6xQX"
},
"source": [
"## Run inference with TF Lite model"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5xr0idac6xfi"
},
"source": [
"Now your TFLite model can be deployed and run using any of the supported [inferencing libraries](https://www.tensorflow.org/lite/guide/inference) or with the new [TFLite AudioClassifier Task API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/audio_classifier). The following code shows how you can run inference with the `.tflite` model in Python."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "nR5zV53YbCIQ"
},
"outputs": [],
"source": [
"# This library provides the TFLite metadata API\n",
"! pip install -q tflite_support"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1AC7PRyiayU5"
},
"outputs": [],
"source": [
"from tflite_support import metadata\n",
"import json\n",
"\n",
"def get_labels(model):\n",
" \"\"\"Returns a list of labels, extracted from the model metadata.\"\"\"\n",
" displayer = metadata.MetadataDisplayer.with_model_file(model)\n",
" labels_file = displayer.get_packed_associated_file_list()[0]\n",
" labels = displayer.get_associated_file_buffer(labels_file).decode()\n",
" return [line for line in labels.split('\\n')]\n",
"\n",
"def get_input_sample_rate(model):\n",
" \"\"\"Returns the model's expected sample rate, from the model metadata.\"\"\"\n",
" displayer = metadata.MetadataDisplayer.with_model_file(model)\n",
" metadata_json = json.loads(displayer.get_metadata_json())\n",
" input_tensor_metadata = metadata_json['subgraph_metadata'][0][\n",
" 'input_tensor_metadata'][0]\n",
" input_content_props = input_tensor_metadata['content']['content_properties']\n",
" return input_content_props['sample_rate']"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yC7TEvQ9o4mu"
},
"source": [
"To observe how well the model performs with real samples, run the following code block over and over. Each time, it will fetch a new test sample and run inference with it, and you can listen to the audio sample below."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "loU6PleipSPf"
},
"outputs": [],
"source": [
"# Get a WAV file for inference and list of labels from the model\n",
"tflite_file = os.path.join(SAVE_PATH, TFLITE_FILENAME)\n",
"labels = get_labels(tflite_file)\n",
"random_audio = get_random_audio_file(test_dir)\n",
"\n",
"# Ensure the audio sample fits the model input\n",
"interpreter = tf.lite.Interpreter(tflite_file)\n",
"input_details = interpreter.get_input_details()\n",
"output_details = interpreter.get_output_details()\n",
"input_size = input_details[0]['shape'][1]\n",
"sample_rate = get_input_sample_rate(tflite_file)\n",
"audio_data, _ = librosa.load(random_audio, sr=sample_rate)\n",
"if len(audio_data) < input_size:\n",
" audio_data.resize(input_size)\n",
"audio_data = np.expand_dims(audio_data[:input_size], axis=0)\n",
"\n",
"# Run inference\n",
"interpreter.allocate_tensors()\n",
"interpreter.set_tensor(input_details[0]['index'], audio_data)\n",
"interpreter.invoke()\n",
"output_data = interpreter.get_tensor(output_details[0]['index'])\n",
"\n",
"# Display prediction and ground truth\n",
"top_index = np.argmax(output_data[0])\n",
"label = labels[top_index]\n",
"score = output_data[0][top_index]\n",
"print('---prediction---')\n",
"print(f'Class: {label}\\nScore: {score}')\n",
"print('----truth----')\n",
"show_sample(random_audio)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VtmfoJW6G2fd"
},
"source": [
"## Download the TF Lite model"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8zLDeiQ_z2Vj"
},
"source": [
"Now you can deploy the TF Lite model to your mobile or embedded device. You don't need to download the labels file because you can instead retrieve the labels from `.tflite` file metadata, as shown in the previous inferencing example."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cNuQoqtjG4zu"
},
"outputs": [],
"source": [
"try:\n",
" from google.colab import files\n",
"except ImportError:\n",
" pass\n",
"else:\n",
" files.download(tflite_file)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "iERuGZz4z6rB"
},
"source": [
"Check out our end-to-end example apps that perform inferencing with TFLite audio models on [Android](https://github.com/tensorflow/examples/tree/master/lite/examples/sound_classification/android/) and [iOS](https://github.com/tensorflow/examples/tree/master/lite/examples/sound_classification/ios)."
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "model_maker_speech_recognition.ipynb",
"private_outputs": true,
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,558 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "h2q27gKz1H20"
},
"source": [
"##### Copyright 2024 The AI Edge Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "TUfAcER1oUS6"
},
"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": "Gb7qyhNL1yWt"
},
"source": [
"# Text Searcher with TensorFlow Lite Model Maker"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Fw5Y7snSuG51"
},
"source": [
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://www.tensorflow.org/lite/models/modify/model_maker/text_searcher\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/text_searcher.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/text_searcher.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/models/modify/model_maker/text_searcher.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://tfhub.dev/google/universal-sentence-encoder-lite/2\"><img src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" />See TF Hub model</a>\n",
" </td>\n",
"</table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "c2sdIlXEVPZR"
},
"source": [
"In this colab notebook, you can learn how to use the [TensorFlow Lite Model Maker](https://www.tensorflow.org/lite/models/modify/model_maker) library to create a TFLite Searcher model. You can use a text Searcher model to build Semantic Search or Smart Reply for your app. This type of model lets you take a text query and search for the most related entries in a text dataset, such as a database of web pages. The model returns a list of the smallest distance scoring entries in the dataset, including metadata you specify, such as URL, page title, or other text entry identifiers. After building this, you can deploy it onto devices (e.g. Android) using [Task Library Searcher API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/text_searcher) to run inference with just a few lines of code.\n",
"\n",
"This tutorial leverages CNN/DailyMail dataset as an instance to create the TFLite Searcher model. You can try with your own dataset with the compatible input comma separated value (CSV) format."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "A_aZYF85VaVK"
},
"source": [
"## Text search using Scalable Nearest Neighbor"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "sr3q-gvm3cI8"
},
"source": [
"This tutorial uses the publicly available CNN/DailyMail non-anonymized summarization dataset, which was produced from the [GitHub repo](https://github.com/abisee/cnn-dailymail). This dataset contains over 300k news articles, which makes it a good dataset to build the Searcher model, and return various related news during model inference for a text query.\n",
"\n",
"The text Searcher model in this example uses a [ScaNN](https://github.com/google-research/google-research/tree/master/scann) (Scalable Nearest Neighbors) index file that can search for similar items from a predefined database. ScaNN achieves state-of-the-art performance for efficient vector similarity search at scale.\n",
"\n",
"Highlights and urls in this dataset are used in this colab to create the model:\n",
"\n",
"1. Highlights are the text for generating the embedding feature vectors and then used for search.\n",
"2. Urls are the returned result shown to users after searching the related highlights.\n",
"\n",
"This tutorial saves these data into the CSV file and then uses the CSV file to build the model. Here are several examples from the dataset.\n",
"\n",
"\n",
"| Highlights | Urls\n",
"| ---------- |----------\n",
"|Hawaiian Airlines again lands at No. 1 in on-time performance. The Airline Quality Rankings Report looks at the 14 largest U.S. airlines. ExpressJet <br> and American Airlines had the worst on-time performance. Virgin America had the best baggage handling; Southwest had lowest complaint rate. | http://www.cnn.com/2013/04/08/travel/airline-quality-report\n",
"| European football's governing body reveals list of countries bidding to host 2020 finals. The 60th anniversary edition of the finals will be hosted by 13 <br> countries. Thirty-two countries are considering bids to host 2020 matches. UEFA will announce host cities on September 25. | http://edition.cnn.com:80/2013/09/20/sport/football/football-euro-2020-bid-countries/index.html?\n",
"| Once octopus-hunter Dylan Mayer has now also signed a petition of 5,000 divers banning their hunt at Seacrest Park. Decision by Washington <br> Department of Fish and Wildlife could take months. | http://www.dailymail.co.uk:80/news/article-2238423/Dylan-Mayer-Washington-considers-ban-Octopus-hunting-diver-caught-ate-Puget-Sound.html?\n",
"| Galaxy was observed 420 million years after the Big Bang. found by NASAs Hubble Space Telescope, Spitzer Space Telescope, and one of natures <br> own natural 'zoom lenses' in space. | http://www.dailymail.co.uk/sciencetech/article-2233883/The-furthest-object-seen-Record-breaking-image-shows-galaxy-13-3-BILLION-light-years-Earth.html\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bcLF2PKkSbV3"
},
"source": [
"## Setup\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2vvAObmTqglq"
},
"source": [
"Start by installing the required packages, including the Model Maker package from the [GitHub repo](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "qhl8lqVamEty"
},
"outputs": [],
"source": [
"!sudo apt -y install libportaudio2\n",
"!pip install -q tflite-model-maker\n",
"!pip install gdown"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "l6lRhVK9Q_0U"
},
"source": [
"Import the required packages."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "XtxiUeZEiXpt"
},
"outputs": [],
"source": [
"from tflite_model_maker import searcher"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "veelLw_VT6uQ"
},
"source": [
"### Prepare the dataset\n",
"\n",
"This tutorial uses the dataset CNN / Daily Mail summarization dataset from the [GitHub repo](https://github.com/abisee/cnn-dailymail).\n",
"\n",
"First, download the text and urls of cnn and dailymail and unzip them. If it\n",
"failed to download from google drive, please wait a few minutes to try it again or download it manually and then upload it to the colab."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "-P3zxue1T6Iy"
},
"outputs": [],
"source": [
"!gdown https://drive.google.com/uc?id=0BwmD_VLjROrfTHk4NFg2SndKcjQ\n",
"!gdown https://drive.google.com/uc?id=0BwmD_VLjROrfM1BxdkxVaTY2bWs\n",
"\n",
"!wget -O all_train.txt https://raw.githubusercontent.com/abisee/cnn-dailymail/master/url_lists/all_train.txt\n",
"!tar xzf cnn_stories.tgz\n",
"!tar xzf dailymail_stories.tgz"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aoOWzTU7ViPM"
},
"source": [
"Then, save the data into the CSV file that can be loaded into `tflite_model_maker` library. The code is based on the logic used to load this data in [`tensorflow_datasets`](https://github.com/tensorflow/datasets/blob/master/tensorflow_datasets/summarization/cnn_dailymail.py). We can't use `tensorflow_dataset` directly since it doesn't contain urls which are used in this colab.\n",
"\n",
"Since it takes a long time to process the data into embedding feature vectors\n",
"for the whole dataset. Only first 5% stories of CNN and Daily Mail dataset are\n",
"selected by default for demo purpose. You can adjust the\n",
"fraction or try with the pre-built TFLite [model](https://storage.googleapis.com/download.tensorflow.org/models/tflite_support/searcher/text_to_image_blogpost/cnn_daily_text_searcher.tflite) with 50% stories of CNN and Daily Mail dataset to search as well."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "bA4PsR6NVU69"
},
"outputs": [],
"source": [
"#@title Save the highlights and urls to the CSV file\n",
"#@markdown Load the highlights from the stories of CNN / Daily Mail, map urls with highlights, and save them to the CSV file.\n",
"\n",
"CNN_FRACTION = 0.05 #@param {type:\"number\"}\n",
"DAILYMAIL_FRACTION = 0.05 #@param {type:\"number\"}\n",
"\n",
"import csv\n",
"import hashlib\n",
"import os\n",
"import tensorflow as tf\n",
"\n",
"dm_single_close_quote = u\"\\u2019\" # unicode\n",
"dm_double_close_quote = u\"\\u201d\"\n",
"END_TOKENS = [\n",
" \".\", \"!\", \"?\", \"...\", \"'\", \"`\", '\"', dm_single_close_quote,\n",
" dm_double_close_quote, \")\"\n",
"] # acceptable ways to end a sentence\n",
"\n",
"\n",
"def read_file(file_path):\n",
" \"\"\"Reads lines in the file.\"\"\"\n",
" lines = []\n",
" with tf.io.gfile.GFile(file_path, \"r\") as f:\n",
" for line in f:\n",
" lines.append(line.strip())\n",
" return lines\n",
"\n",
"\n",
"def url_hash(url):\n",
" \"\"\"Gets the hash value of the url.\"\"\"\n",
" h = hashlib.sha1()\n",
" url = url.encode(\"utf-8\")\n",
" h.update(url)\n",
" return h.hexdigest()\n",
"\n",
"\n",
"def get_url_hashes_dict(urls_path):\n",
" \"\"\"Gets hashes dict that maps the hash value to the original url in file.\"\"\"\n",
" urls = read_file(urls_path)\n",
" return {url_hash(url): url[url.find(\"id_/\") + 4:] for url in urls}\n",
"\n",
"\n",
"def find_files(folder, url_dict):\n",
" \"\"\"Finds files corresponding to the urls in the folder.\"\"\"\n",
" all_files = tf.io.gfile.listdir(folder)\n",
" ret_files = []\n",
" for file in all_files:\n",
" # Gets the file name without extension.\n",
" filename = os.path.splitext(os.path.basename(file))[0]\n",
" if filename in url_dict:\n",
" ret_files.append(os.path.join(folder, file))\n",
" return ret_files\n",
"\n",
"\n",
"def fix_missing_period(line):\n",
" \"\"\"Adds a period to a line that is missing a period.\"\"\"\n",
" if \"@highlight\" in line:\n",
" return line\n",
" if not line:\n",
" return line\n",
" if line[-1] in END_TOKENS:\n",
" return line\n",
" return line + \".\"\n",
"\n",
"\n",
"def get_highlights(story_file):\n",
" \"\"\"Gets highlights from a story file path.\"\"\"\n",
" lines = read_file(story_file)\n",
"\n",
" # Put periods on the ends of lines that are missing them\n",
" # (this is a problem in the dataset because many image captions don't end in\n",
" # periods; consequently they end up in the body of the article as run-on\n",
" # sentences)\n",
" lines = [fix_missing_period(line) for line in lines]\n",
"\n",
" # Separate out article and abstract sentences\n",
" highlight_list = []\n",
" next_is_highlight = False\n",
" for line in lines:\n",
" if not line:\n",
" continue # empty line\n",
" elif line.startswith(\"@highlight\"):\n",
" next_is_highlight = True\n",
" elif next_is_highlight:\n",
" highlight_list.append(line)\n",
"\n",
" # Make highlights into a single string.\n",
" highlights = \"\\n\".join(highlight_list)\n",
"\n",
" return highlights\n",
"\n",
"url_hashes_dict = get_url_hashes_dict(\"all_train.txt\")\n",
"cnn_files = find_files(\"cnn/stories\", url_hashes_dict)\n",
"dailymail_files = find_files(\"dailymail/stories\", url_hashes_dict)\n",
"\n",
"# The size to be selected.\n",
"cnn_size = int(CNN_FRACTION * len(cnn_files))\n",
"dailymail_size = int(DAILYMAIL_FRACTION * len(dailymail_files))\n",
"print(\"CNN size: %d\"%cnn_size)\n",
"print(\"Daily Mail size: %d\"%dailymail_size)\n",
"\n",
"with open(\"cnn_dailymail.csv\", \"w\") as csvfile:\n",
" writer = csv.DictWriter(csvfile, fieldnames=[\"highlights\", \"urls\"])\n",
" writer.writeheader()\n",
"\n",
" for file in cnn_files[:cnn_size] + dailymail_files[:dailymail_size]:\n",
" highlights = get_highlights(file)\n",
" # Gets the filename which is the hash value of the url.\n",
" filename = os.path.splitext(os.path.basename(file))[0]\n",
" url = url_hashes_dict[filename]\n",
" writer.writerow({\"highlights\": highlights, \"urls\": url})\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xushUyZXqP59"
},
"source": [
"## Build the text Searcher model"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vn61LJ9QbOPi"
},
"source": [
"Create a text Searcher model by loading a dataset, creating a model with the data and exporting the TFLite model.\n",
"\n",
"### Step 1. Load the dataset\n",
"\n",
"Model Maker takes the text dataset and the corresponding metadata of each text string (such as urls in this example) in the CSV format. It embeds the text strings into feature vectors using the user-specified embedder model.\n",
"\n",
"In this demo, we build the Searcher model using [Universal Sentence Encoder](https://tfhub.dev/google/universal-sentence-encoder-lite/2), a state-of-the-art sentence embedding model which is already retrained from [colab](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/examples/colab/on_device_text_to_image_search_tflite.ipynb). The model is optimized for on-device inference performance, and only takes 6ms to embed a query string (measured on Pixel 6). Alternatively, you can use [this](https://tfhub.dev/google/lite-model/universal-sentence-encoder-qa-ondevice/1?lite-format=tflite) quantized version, which is smaller but takes 38ms for each embedding."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1ymHbk0wjHHZ"
},
"outputs": [],
"source": [
"!wget -O universal_sentence_encoder.tflite https://storage.googleapis.com/download.tensorflow.org/models/tflite_support/searcher/text_to_image_blogpost/text_embedder.tflite"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "knJwmJHxkFbx"
},
"source": [
"Create a `searcher.TextDataLoader` instance and use `data_loader.load_from_csv` method to load the dataset. It takes ~10 minutes for this\n",
"step since it generates the embedding feature vector for each text one by one. You can try to upload your own CSV file and load it to build the customized model as well.\n",
"\n",
"Specify the name of text column and metadata column in the CSV file.\n",
"* Text is used to generate the embedding feature vectors.\n",
"* Metadata is the content to be shown when you search the certain text.\n",
"\n",
"Here are the first 4 lines of the CNN-DailyMail CSV file generated above.\n",
"\n",
"| highlights| urls\n",
"| ---------- |----------\n",
"|Syrian official: Obama climbed to the top of the tree, doesn't know how to get down. Obama sends a letter to the heads of the House and Senate. Obama <br> to seek congressional approval on military action against Syria. Aim is to determine whether CW were used, not by whom, says U.N. spokesman.|http://www.cnn.com/2013/08/31/world/meast/syria-civil-war/\n",
"|Usain Bolt wins third gold of world championship. Anchors Jamaica to 4x100m relay victory. Eighth gold at the championships for Bolt. Jamaica double <br> up in women's 4x100m relay.|http://edition.cnn.com/2013/08/18/sport/athletics-bolt-jamaica-gold\n",
"|The employee in agency's Kansas City office is among hundreds of \"virtual\" workers. The employee's travel to and from the mainland U.S. last year cost <br> more than $24,000. The telecommuting program, like all GSA practices, is under review.|http://www.cnn.com:80/2012/08/23/politics/gsa-hawaii-teleworking\n",
"|NEW: A Canadian doctor says she was part of a team examining Harry Burkhart in 2010. NEW: Diagnosis: \"autism, severe anxiety, post-traumatic stress <br> disorder and depression\" Burkhart is also suspected in a German arson probe, officials say. Prosecutors believe the German national set a string of fires <br> in Los Angeles.|http://edition.cnn.com:80/2012/01/05/justice/california-arson/index.html?\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "CtdZ-JDwMimd"
},
"outputs": [],
"source": [
"data_loader = searcher.TextDataLoader.create(\"universal_sentence_encoder.tflite\", l2_normalize=True)\n",
"data_loader.load_from_csv(\"cnn_dailymail.csv\", text_column=\"highlights\", metadata_column=\"urls\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VVN5bkSFiZdV"
},
"source": [
"For image use cases, you can create a `searcher.ImageDataLoader` instance and then use `data_loader.load_from_folder` to load images from the folder. The `searcher.ImageDataLoader` instance needs to be created by a TFLite embedder model because it will be leveraged to encode queries to feature vectors and be exported with the TFLite Searcher model. For instance:\n",
"```python\n",
"data_loader = searcher.ImageDataLoader.create(\"mobilenet_v2_035_96_embedder_with_metadata.tflite\")\n",
"data_loader.load_from_folder(\"food/\")\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2uZkLR6N6gDR"
},
"source": [
"###Step 2. Create the Searcher model\n",
"\n",
"* Configure ScaNN options. See [api doc](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/searcher/ScaNNOptions) for more details.\n",
"* Create the Searcher model from data and ScaNN options. You can see the [in-depth examination](https://ai.googleblog.com/2020/07/announcing-scann-efficient-vector.html) to learn more about the ScaNN algorithm."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "kwlYdTcg63xy"
},
"outputs": [],
"source": [
"scann_options = searcher.ScaNNOptions(\n",
" distance_measure=\"dot_product\",\n",
" tree=searcher.Tree(num_leaves=140, num_leaves_to_search=4),\n",
" score_ah=searcher.ScoreAH(dimensions_per_block=1, anisotropic_quantization_threshold=0.2))\n",
"model = searcher.Searcher.create_from_data(data_loader, scann_options)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "lluAWms3soFm"
},
"source": [
"In the above example, we define the following options:\n",
"* `distance_measure`: we use \"dot_product\" to measure the distance between two embedding vectors. Note that we actually compute the **negative** dot product value to preserve the notion that \"smaller is closer\".\n",
"\n",
"* `tree`: the dataset is divided the dataset into 140 partitions (roughly the square root of the data size), and 4 of them are searched during retrieval, which is roughly 3% of the dataset.\n",
"\n",
"* `score_ah`: we quantize the float embeddings to int8 values with the same dimension to save space."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "CgCDMe0e6jlT"
},
"source": [
"###Step 3. Export the TFLite model\n",
"\n",
"Then you can export the TFLite Searcher model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Hm_UULdW7A9T"
},
"outputs": [],
"source": [
"model.export(\n",
" export_filename=\"searcher.tflite\",\n",
" userinfo=\"\",\n",
" export_format=searcher.ExportFormat.TFLITE)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "me6_RwPZqNhX"
},
"source": [
"## Test the TFLite model on your query\n",
"\n",
"You can test the exported TFLite model using custom query text. To query text using the Searcher model, initialize the model and run a search with text phrase, as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "GkXtipXKqXp4"
},
"outputs": [],
"source": [
"from tflite_support.task import text\n",
"\n",
"# Initializes a TextSearcher object.\n",
"searcher = text.TextSearcher.create_from_file(\"searcher.tflite\")\n",
"\n",
"# Searches the input query.\n",
"results = searcher.search(\"The Airline Quality Rankings Report looks at the 14 largest U.S. airlines.\")\n",
"print(results)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GzI0tPM2rFlc"
},
"source": [
"See the [Task Library documentation](https://www.tensorflow.org/lite/inference_with_metadata/task_library/text_searcher) for more information about how to integrate the model to various platforms."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HS4u77W5gnzQ"
},
"source": [
"# Read more\n",
"\n",
" For more information, please refer to:\n",
"\n",
"* TensorFlow Lite Model Maker [guide](https://www.tensorflow.org/lite/models/modify/model_maker) and [API reference](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker).\n",
"\n",
"* Task Library: [TextSearcher](https://www.tensorflow.org/lite/inference_with_metadata/task_library/text_searcher) for deployment.\n",
"* The end-to-end reference apps: [Android](https://github.com/tensorflow/examples/tree/master/lite/examples/text_searcher/android).\n"
]
}
],
"metadata": {
"colab": {
"collapsed_sections": [],
"name": "Model Maker Text Searcher Tutorial",
"provenance": [
{
"file_id": "1dbRXQCjtm-jBFC32DJ6YCVXnXBOG3M5t",
"timestamp": 1613441434239
},
{
"file_id": "https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/modify/model_maker/text_classification.ipynb",
"timestamp": 1612303859066
}
],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
@@ -0,0 +1,102 @@
# Pre-trained models for TensorFlow Lite
There are a variety of already trained, open source models you can use
immediately with TensorFlow Lite to accomplish many machine learning tasks.
Using pre-trained TensorFlow Lite models lets you add machine learning
functionality to your mobile and edge device application quickly, without having
to build and train a model. This guide helps you find and decide on trained
models for use with TensorFlow Lite.
You can start browsing TensorFlow Lite models right away based on general use
cases in the [TensorFlow Lite Examples](../../examples) section, or browse a
larger set of models on [TensorFlow Hub](https://tfhub.dev/s?deployment-
format=lite).
**Important:** TensorFlow Hub lists both regular TensorFlow models and
TensorFlow Lite format models. These model formats are not interchangeable.
TensorFlow models can be converted into TensorFlow Lite models, but that process
is not reversible.
## Find a model for your application
Finding an existing TensorFlow Lite model for your use case can be tricky
depending on what you are trying to accomplish. Here are a few recommended ways
to discover models for use with TensorFlow Lite:
**By example:** The fastest way to find and start using models with TensorFlow
Lite is to browse the [TensorFlow Lite Examples](../../examples) section to find
models that perform a task which is similar to your use case. This short catalog
of examples provides models for common use cases with explanations of the models
and sample code to get you started running and using them.
**By data input type:** Aside from looking at examples similar to your use
case, another way to discover models for your own use is to consider the type of
data you want to process, such as audio, text, images, or video data. Machine
learning models are frequently designed for use with one of these types of data,
so looking for models that handle the data type you want to use can help you
narrow down what models to consider. On [TensorFlow
Hub](https://tfhub.dev/s?deployment-format=lite), you can use the **Problem
domain** filter to view model data types and narrow your list.
Note: Processing video with machine learning models can frequently be
accomplished with models that are designed for processing single images,
depending on how fast and how many inferences you need to perform for your use
case. If you intend to use video for your use case, consider using single-frame
video sampling with a model built for fast processing of individual images.
The following lists links to TensorFlow Lite models on [TensorFlow
Hub](https://tfhub.dev/s?deployment-format=lite) for common use cases:
- [Image classification](https://tfhub.dev/s?deployment-format=lite&module-type=image-classification)
models
- [Object detection](https://tfhub.dev/s?deployment-format=lite&module-type=image-object-detection)
models
- [Text classification](https://tfhub.dev/s?deployment-format=lite&module-type=text-classification)
models
- [Text embedding](https://tfhub.dev/s?deployment-format=lite&module-type=text-embedding)
models
- [Audio speech synthesis](https://tfhub.dev/s?deployment-format=lite&module-type=audio-speech-synthesis)
models
- [Audio embedding](https://tfhub.dev/s?deployment-format=lite&module-type=audio-embedding)
models
## Choose between similar models
If your application follows a common use case such as image classification or
object detection, you may find yourself deciding between multiple TensorFlow
Lite models, with varying binary size, data input size, inference speed, and
prediction accuracy ratings. When deciding between a number of models, you
should narrow your options based first on your most limiting constraint: size of
model, size of data, inference speed, or accuracy.
Key Point: Generally, when choosing between similar models, pick the smallest
model to allow for the broadest device compatibility and fast inference times.
If you are not sure what your most limiting constraint is, assume it is the
size of the model and pick the smallest model available. Picking a small model
gives you the most flexibility in terms of the devices where you can
successfully deploy and run the model. Smaller models also typically produce
faster inferences, and speedier predictions generally create better end-user
experiences. Smaller models typically have lower accuracy rates, so you may need
to pick larger models if prediction accuracy is your primary concern.
## Sources for models
Use the [TensorFlow Lite Examples](../../examples)
section and [TensorFlow Hub](https://tfhub.dev/s?deployment-format=lite) as your
first destinations for finding and selecting models for use with TensorFlow
Lite. These sources generally have up to date, curated models for use with
TensorFlow Lite, and frequently include sample code to accelerate your
development process.
### TensorFlow models
It is possible to [convert](https://www.tensorflow.org/lite/models/convert) regular
TensorFlow models to TensorFlow Lite format. For more information about
converting models, see the [TensorFlow Lite
Converter](https://www.tensorflow.org/lite/models/convert) documentation. You can find
TensorFlow models on [TensorFlow Hub](https://tfhub.dev/) and in the
[TensorFlow Model Garden](https://github.com/tensorflow/models).