Tensorflow Lite 入门——模型的训练和转换

请设置文章作者

发布于：Nov 22, 2020

之前的文章中介绍的使用 Tensorflow 训练的模型要么只能运行在 PC 端，要么需要云端的支持。随着智能手机和物联网设备的普及，能够在智能手机甚至嵌入式设备直接运行的模型需求就越来越高。这篇文章就开始介绍 Tensorflow Lite, 这个能够运行在智能手机和嵌入式设备的开源深度学习框架。

通常，我们会在 PC 或者云端建立模型，并对模型进行训练，然后将模型转换成 Tensorflow Lite 的格式，并最终部署到终端设备上，这篇文章我们就用 Fashion MNIST 的数据集，建立并训练模型，并采用模拟器的方式部署到终端设备上进行测试。

1. 数据的加载和训练

这个部分的内容与之前Tensorflow 2.0 快速入门内容重复，在这里就不过多赘述了。但是值得注意的是，之前的文章我们都是使用的 keras 的数据集，其数据格式是 numpy。Fasion MINST 是从 tensorflow_datasets 中直接加载的。将数据集分为 80% 训练集，10% 测试集和10%验证集。tfds的详细使用说明请参考官方文档。

import tensorflow_datasets as tfds

splits, info = tfds.load('fashion_mnist', with_info=True, as_supervised=True, 
                         split=['train[:80%]', 'train[80%:90%]', 'train[90%:]'])

(train_examples, validation_examples, test_examples) = splits

num_examples = info.splits['train'].num_examples
num_classes = info.features['label'].num_classes

数据经过预处理之后，我们使用 Keras 的 API 快速搭建了一个五层的 CNN 神经网络并使用 model.fit 对模型进行了训练。

2. 模型的保存与转换

训练好的模型这里使用了 tf.save_model.save() 将模型保存在了指定目录。

export_dir = 'saved_model/1'

# Use the tf.saved_model API to export the SavedModel
tf.saved_model.save(model, export_dir)

本文的重点是模型转换，在Tensorflow Lite 中，使用 TFLiteCoverter 可以轻松将模型转换成 Tensorflow Lite 的模型。

# Use the TFLiteConverter SavedModel API to initialize the converter
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir=export_dir)

# Invoke the converter to finally generate the TFLite model
tflite_model = converter.convert()

# Save the model file as 'model.tflite'
tflite_model_file = 'model.tflite'
with open(tflite_model_file, "wb") as f:
  f.write(tflite_model)

3. 模型的优化

converter 在默认的情况下是将模型权重从32位浮点数转换成8位整数从而大大减小模型的大小。

1	converter.optimizations = [tf.lite.Optimize.DEFAULT]

我们也可以将模型手动调整为16位浮点数

1 2	converter.optimizations = [tf.lite.Optimize.DEFAULT] converter.target_spec.supported_types = [tf.float16]

当然除了 default 模式，官方也提供了精度优先和大小优先的优化模式，详细内容参考官方文档

1 2	tf.lite.Optimize.OPTIMIZE_FOR_LATENCY tf.lite.Optimize.OPTIMIZE_FOR_SIZE

4. 模型测试

到这里，实际上我们就可以将转换成 Tensorflow Lite 的模型部署到设备上进行测试了，但是此时，我们并不知道模型的性能如何，Tensorflow Lite 提供了模拟器，我们可以轻松部署在模拟器上对转换后的模型进行测试。

测试模型分为三步，

第一步：加载 TFLite 模型，部署tensor

第二步：获取 input 和 output index

第三步：加载数据并获取结果

# Step 1: Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

# Step 2: Get input and output tensors index
input_index = interpreter.get_input_details()[0]["index"]
output_index = interpreter.get_output_details()[0]["index"]

# Step 3: Get results
interpreter.set_tensor(input_index, img)
interpreter.invoke()
predictions = interpreter.get_tensor(output_index)