lobe/lobe-python

Feature Request: Export/Import trained models as Keras HDF5 format

PedramKhalaj opened this issue · 1 comments

Currently, the lobe desktop application only exports the trained models as TensorFlow 1 data format. However, it would be nice to have an option to export the models as the Keras HDF5 format.

I have trained a model using lobe desktop application and want to save/load it as a Keras HDF5 format. I have tried the suggestion tensorflow/tensorflow#42425 without any luck.

Any comments on this from the team? Any workarounds?

From some searching, it looks difficult to do. We save to a TF 1.15 SavedModel, but that loads fine in TF 2.X (which this library uses).

Here is a full example script I got working with TF 2.6 and our SavedModel export, run it in the same directory as the Lobe export:

import tensorflow as tf
import numpy as np
import lobe


loaded = tf.saved_model.load(".")
sig = lobe.signature.ImageClassificationSignature(".")
input_shape = sig.input_image_size + (3,)


class LayerFromSavedModel(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(LayerFromSavedModel, self).__init__(**kwargs)
        self.vars = loaded.variables

    def call(self, inputs):
        return loaded.signatures["serving_default"](inputs)


input = tf.keras.Input(shape=input_shape, dtype=np.float32)
model = tf.keras.Model(input, LayerFromSavedModel()(input))
model.save("keras_model.h5")


model2 = tf.keras.models.load_model(
    "keras_model.h5", custom_objects={"LayerFromSavedModel": LayerFromSavedModel}
)
print(model2.summary())

adapted from the tf issue here: tensorflow/tensorflow#42425