##### Copyright 2021 The TensorFlow Authors.
#@title 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
#
# https://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.
将您的 TFLite 代码迁移到 TF2#
 在 TensorFlow.org 上查看 |
 在 Google Colab 中运行 |
 在 GitHub 上查看源代码 |
 下载笔记本 |
TensorFlow Lite (TFLite) 是一套工具，可帮助开发者在设备端（移动、嵌入式和物联网设备）上运行机器学习推断。TFLite 转换器可以将现有 TF 模型转换为可在设备端高效运行的优化 TFLite 模型格式。
在本文档中，您将了解需要对 TF 到 TFLite 的转换代码进行哪些更改，然后是几个实现相同目标的示例。
TF 到 TFLite 转换代码的更改#
如果您使用的是旧版 TF1 模型格式（例如，Keras 文件、冻结的 GraphDef、检查点、tf.Session 等），请将其更新为 TF1/TF2 SavedModel，并使用 TF2 转换器 API
tf.lite.TFLiteConverter.from_saved_model(...)将其转换为 TFLite 模型（请参见表 1）。更新转换器 API 标志（请参见表 2）。
移除旧版 API，例如
tf.lite.constants。（例如：将tf.lite.constants.INT8替换为tf.int8）
// 表 1 // TFLite Python 转换器 API 更新
TF1 API |
TF2 API |
|---|---|
|
支持 |
|
已移除（更新为 SavedModel 格式） |
|
已移除（更新为 SavedModel 格式） |
|
已移除（更新为 SavedModel 格式） |
<style> .table {margin-left: 0 !important;} </style>
// 表 2 // TFLite Python 转换器 API 标志更新
TF1 API |
TF2 API |
|---|---|
|
支持 |
|
已移除（不支持的转换器 API 参数） |
|
已移除（不支持的量化工作流） |
|
已移除（改为使用 Netron 或 visualize.py 呈现模型） |
|
已移除（TF2 中不支持的功能） |
示例#
您现在将演练一些示例，将旧版 TF1 模型转换为 TF1/TF2 SavedModel，然后将其转换为 TF2 TFLite 模型。
安装#
从必要的 TensorFlow 导入开始。
import tensorflow as tf
import tensorflow.compat.v1 as tf1
import numpy as np
import logging
logger = tf.get_logger()
logger.setLevel(logging.ERROR)
import shutil
def remove_dir(path):
try:
shutil.rmtree(path)
except:
pass
创建所有必要的 TF1 模型格式。
# Create a TF1 SavedModel
SAVED_MODEL_DIR = "tf_saved_model/"
remove_dir(SAVED_MODEL_DIR)
with tf1.Graph().as_default() as g:
with tf1.Session() as sess:
input = tf1.placeholder(tf.float32, shape=(3,), name='input')
output = input + 2
# print("result: ", sess.run(output, {input: [0., 2., 4.]}))
tf1.saved_model.simple_save(
sess, SAVED_MODEL_DIR,
inputs={'input': input},
outputs={'output': output})
print("TF1 SavedModel path: ", SAVED_MODEL_DIR)
# Create a TF1 Keras model
KERAS_MODEL_PATH = 'tf_keras_model.h5'
model = tf1.keras.models.Sequential([
tf1.keras.layers.InputLayer(input_shape=(128, 128, 3,), name='input'),
tf1.keras.layers.Dense(units=16, input_shape=(128, 128, 3,), activation='relu'),
tf1.keras.layers.Dense(units=1, name='output')
])
model.save(KERAS_MODEL_PATH, save_format='h5')
print("TF1 Keras Model path: ", KERAS_MODEL_PATH)
# Create a TF1 frozen GraphDef model
GRAPH_DEF_MODEL_PATH = tf.keras.utils.get_file(
'mobilenet_v1_0.25_128',
origin='https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_0.25_128_frozen.tgz',
untar=True,
) + '/frozen_graph.pb'
print("TF1 frozen GraphDef path: ", GRAPH_DEF_MODEL_PATH)
1. 将 TF1 SavedModel 转换为 TFLite 模型#
之前：使用 TF1 进行转换#
下面是 TF1 样式 TFlite 转换的典型代码。
converter = tf1.lite.TFLiteConverter.from_saved_model(
saved_model_dir=SAVED_MODEL_DIR,
input_arrays=['input'],
input_shapes={'input' : [3]}
)
converter.optimizations = {tf.lite.Optimize.DEFAULT}
converter.change_concat_input_ranges = True
tflite_model = converter.convert()
# Ignore warning: "Use '@tf.function' or '@defun' to decorate the function."
之后：使用 TF2 进行转换#
将 TF1 SavedModel 直接转换为 TFLite 模型，并设置较小的 v2 转换器标志。
# Convert TF1 SavedModel to a TFLite model.
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir=SAVED_MODEL_DIR)
converter.optimizations = {tf.lite.Optimize.DEFAULT}
tflite_model = converter.convert()
2. 将 TF1 Keras 模型文件转换为 TFLite 模型#
之前：使用 TF1 进行转换#
下面是 TF1 样式 TFlite 转换的典型代码。
converter = tf1.lite.TFLiteConverter.from_keras_model_file(model_file=KERAS_MODEL_PATH)
converter.optimizations = {tf.lite.Optimize.DEFAULT}
converter.change_concat_input_ranges = True
tflite_model = converter.convert()
之后：使用 TF2 进行转换#
首先，将 TF1 Keras 模型文件转换为 TF2 SavedModel，然后将其转换为 TFLite 模型，并设置较小的 v2 转换器标志。
# Convert TF1 Keras model file to TF2 SavedModel.
model = tf.keras.models.load_model(KERAS_MODEL_PATH)
model.save(filepath='saved_model_2/')
# Convert TF2 SavedModel to a TFLite model.
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir='saved_model_2/')
tflite_model = converter.convert()
3. 将 TF1 冻结的 GraphDef 转换为 TFLite 模型#
之前：使用 TF1 进行转换#
下面是 TF1 样式 TFlite 转换的典型代码。
converter = tf1.lite.TFLiteConverter.from_frozen_graph(
graph_def_file=GRAPH_DEF_MODEL_PATH,
input_arrays=['input'],
input_shapes={'input' : [1, 128, 128, 3]},
output_arrays=['MobilenetV1/Predictions/Softmax'],
)
converter.optimizations = {tf.lite.Optimize.DEFAULT}
converter.change_concat_input_ranges = True
tflite_model = converter.convert()
之后：使用 TF2 进行转换#
首先，将 TF1 冻结的 GraphDef 转换为 TF1 SavedModel，然后将其转换为 TFLite 模型，并设置较小的 v2 转换器标志。
## Convert TF1 frozen Graph to TF1 SavedModel.
# Load the graph as a v1.GraphDef
import pathlib
gdef = tf.compat.v1.GraphDef()
gdef.ParseFromString(pathlib.Path(GRAPH_DEF_MODEL_PATH).read_bytes())
# Convert the GraphDef to a tf.Graph
with tf.Graph().as_default() as g:
tf.graph_util.import_graph_def(gdef, name="")
# Look up the input and output tensors.
input_tensor = g.get_tensor_by_name('input:0')
output_tensor = g.get_tensor_by_name('MobilenetV1/Predictions/Softmax:0')
# Save the graph as a TF1 Savedmodel
remove_dir('saved_model_3/')
with tf.compat.v1.Session(graph=g) as s:
tf.compat.v1.saved_model.simple_save(
session=s,
export_dir='saved_model_3/',
inputs={'input':input_tensor},
outputs={'output':output_tensor})
# Convert TF1 SavedModel to a TFLite model.
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir='saved_model_3/')
converter.optimizations = {tf.lite.Optimize.DEFAULT}
tflite_model = converter.convert()
延伸阅读#
请参阅 TFLite 指南来详细了解工作流和最新功能。
如果您使用的是 TF1 代码或旧版 TF1 模型格式（Keras
.h5文件、冻结的 GraphDef.pb等），请更新您的代码并将您的模型迁移到 TF2 SavedModel 格式。