Tensorflow-Lite转换自带算子的Keras模型(以EfficientNet为例)

谷歌的深度学习框架一贯宣传其工业部署上的领先,代表产品就是Tensorflow旗下的Tensorflow-Lite。

TensorFlow Litetensorflow.google.cn图标

不过,Tensorflow-Lite一般存在几个问题:

  1. Tensroflow版本众多,API接口不稳定,导致Lite的API也不稳定。
  2. 谷歌喜欢发时髦论文,但Lite一般只支持在工业界沉淀下来的算法模型。

不过,WorkAround依然是有的。

我们以他们今年引以为豪的EfficientNet系列模型为例。

这里不涉及到论文细节和代码实现细节,各位看官请移步这里看论文和代码:

https://arxiv.org/abs/1905.11946arxiv.org
qubvel/efficientnetgithub.com图标

有厉害的模型,但怎么部署到轻量级设备上呢?

a. 请一定要装tensorflow 2.0。

b. 如果您的模型是在原生Keras训练的,在转换轻量级模型的时候,请把tf.keras当keras(从2.0开始,谷歌把Keras集成到Tensorflow里,打算跟Pytorch死磕啦)。

c. 实现算子的一些API要拷贝一遍到本地。

必要的库的引入:

import tensorflow as tf
import tensorflow.keras as keras
import h5py
import os
from efficientnet import model
import functools
from tensorflow.python.keras.utils import CustomObjectScope,get_custom_objects

算子要重新定义:

def inject_keras_modules(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        kwargs['backend'] = keras.backend
        kwargs['layers'] = keras.layers
        kwargs['models'] = keras.models
        kwargs['utils'] = keras.utils
        return func(*args, **kwargs)

    return wrapper
def init_keras_custom_objects():
    custom_objects = {
        'swish': inject_keras_modules(model.get_swish)(),
        'FixedDropout': inject_keras_modules(model.get_dropout)()
    }

    get_custom_objects().update(custom_objects)

主程序:比较麻烦,先是用tf.keras加载模型,再存成saved_model格式,同时初始化必要的算子,参考这里:

模型转换器(Converter)的 Python API 指南 | TensorFlow Litetensorflow.google.cn
init_keras_custom_objects()
keras_model_name = 'efficient_net_b0.h5'
keras_model_path = os.path.join('keras_models', keras_model_name)
save_model = tf.keras.models.load_model(keras_model_path)
export_dir='save'
tf.saved_model.save(save_model, export_dir)
new_model = tf.saved_model.load(export_dir)

在定义算子的Scope范围内进行转换:

with CustomObjectScope({'swish': inject_keras_modules(model.get_swish)(),
                        'FixedDropout': inject_keras_modules(model.get_dropout)()}):
    concrete_func = new_model.signatures[
        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    concrete_func.inputs[0].set_shape([1, 64, 64, 3])
    converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])

注意,这里需要声明input-shape,不然会报error-message:

ValueError: None is only supported in the 1st dimension. Tensor 'input_1' has invalid shape '[None, None, None, 3]'.
   concrete_func = new_model.signatures[
        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    concrete_func.inputs[0].set_shape([1, 64, 64, 3])
    converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])

然后,就可以愉快地转换了。

converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
converter.allow_custom_ops = True
tflite_model = converter.convert()
open("efficient_net_b0.tflite", "wb").write(tflite_model)

结论:整个流程还是有些好事多磨,一般来说,还是建议大家多在数据端花功夫,时髦模型可以尝试,比如带ImageNet权重的一些State Of The Art模型,但精度上的提升可能赶不上看官多做些工程上的细活。

发布于 2019-11-01

文章被以下专栏收录