如何把Tensorflow模型转换成TFLite模型

深度学习迅猛发展，目前已经可以移植到移动端使用了，TensorFlow推出的TensorFlow Lite就是一款把深度学习应用到移动端的框架技术。

使用TensorFlowLite 需要tflite文件模型，这个模型可以由TensorFlow训练的模型转换而成。所以首先需要知道如何保存训练好的TensorFlow模型。

一般有这几种保存形式：

1. Checkpoints
2. HDF5
3. SavedModel等

保存与读取CheckPoint

当模型训练结束，可以用以下代码把权重保存成checkpoint格式

model.save_weights('./MyModel',True)

checkpoints文件仅是保存训练好的权重，不带网络结构，所以做predict时需要结合model使用
如：

model = keras_segmentation.models.segnet.mobilenet_segnet(n_classes=2, input_height=224, input_width=224)
model.load_weights('./MyModel')

保存成H5

把训练好的网络保存成h5文件很简单

model.save('MyModel.h5')

H5转换成TFLite

这里是文章主要内容

我习惯使用H5文件转换成tflite文件

官网代码是这样的

converter = tf.lite.TFLiteConverter.from_keras_model_file('newModel.h5')
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)

但我用的keras 2.2.4版本会报下面错误，好像说是新版的keras把relu6改掉了，找不到方法
ValueError: Unknown activation function:relu6

于是需要自己定义一个relu6

import tensorflow as tf
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.utils import CustomObjectScope

def relu6(x):
  return K.relu(x, max_value=6)

with CustomObjectScope({'relu6': relu6}):
    converter = tf.lite.TFLiteConverter.from_keras_model_file('newModel.h5')
    tflite_model = converter.convert()
    open("newModel.tflite", "wb").write(tflite_model)

看到生成的tflite文件表示保存成功了

也可以这么查看tflite网络的输入输出

import numpy as np
import tensorflow as tf

# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path="newModel.tflite")
interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

print(input_details)
print(output_details)

输出了以下信息

[{'name': 'input_1', 'index': 115, 'shape': array([  1, 224, 224,   3]), 'dtype': <class 'numpy.float32'>, 'quantization': (0.0, 0)}]

[{'name': 'activation_1/truediv', 'index': 6, 'shape': array([    1, 12544,     2]), 'dtype': <class 'numpy.float32'>, 'quantization': (0.0, 0)}]

两个shape分别表示输入输出的numpy数组结构，dtype是数据类型

下一篇会分享如何在PC端测试TFLite模型