tests for savedmodel functionality under tf implementation import os import numpy as np import pytest import tensorflow as tf from keras import backend from keras import layers from keras import metrics from keras import models from keras import optimizers from keras import testing from keras saving import objectregistration objectregistration registerkerasserializablepackagemypackage class custommodelxmodels model def initself args kwargs super initargs kwargs self dense1 layers dense1 self dense2 layers dense1 def callself inputs out self dense1inputs return self dense2out def oneself return 1 objectregistration registerkerasserializablepackagemypackage class customsignaturemodelmodels model def initself supercustomsignaturemodel self init self v tf variable1 0 tf function def callself x return x self v tf functioninputsignaturetf tensorspec tf float32 def mutateself newv self v assignnewv pytest mark skipif backend backend tensorflow reasonthe savedmodel test can only run with tf backend class savedmodeltesttesting testcase def testsequentialself model models sequentiallayers dense1 model compilelossmse optimizeradam xtrain np random rand100 3 ytrain np random rand100 1 model fitxtrain ytrain path os path joinself gettempdir mykerasmodel tf savedmodel savemodel path restoredmodel tf savedmodel loadpath self assertallclose modelxtrain restoredmodel signaturesservingdefault tf converttotensorxtrain dtypetf float32 output0 rtol1e4 atol1e4 def testfunctionalself inputs layers inputshape3 x layers dense1 namefirstdenseinputs outputs layers dense1 nameseconddensex model models modelinputs outputs model compile optimizeradam lossmse xtrain np random rand100 3 ytrain np random rand100 1 model fitxtrain ytrain path os path joinself gettempdir mykerasmodel tf savedmodel savemodel path restoredmodel tf savedmodel loadpath self assertallclose modelxtrain restoredmodel signaturesservingdefault tf converttotensorxtrain dtypetf float32 output0 rtol1e4 atol1e4 def testsubclassedself model custommodelx model compile optimizeradam lossmse metricsmetrics hinge mse xtrain np random rand100 3 ytrain np random rand100 1 model fitxtrain ytrain path os path joinself gettempdir mykerasmodel tf savedmodel savemodel path restoredmodel tf savedmodel loadpath self assertallclose modelxtrain restoredmodel signaturesservingdefault tf converttotensorxtrain dtypetf float32 output0 rtol1e4 atol1e4 def testcustommodelandlayerself objectregistration registerkerasserializablepackagemypackage class customlayerlayers layer def callself inputs return inputs objectregistration registerkerasserializablepackagemypackage class modelmodels model def initself super init self layer customlayer tf functioninputsignaturetf tensorspecnone 1 def callself inputs return self layerinputs model model inp np array1 0 result modelinp path os path joinself gettempdir mykerasmodel tf savedmodel savemodel path restoredmodel tf savedmodel loadpath self assertallclose result restoredmodel callinp rtol1e4 atol1e4 def testmultiinputmodelself input1 layers inputshape3 input2 layers inputshape5 model models modelinput1 input2 input1 input2 path os path joinself gettempdir mykerasmodel tf savedmodel savemodel path restoredmodel tf savedmodel loadpath inputarr1 np random random1 3 astypefloat32 inputarr2 np random random1 5 astypefloat32 outputs restoredmodel signaturesservingdefault inputstf converttotensorinputarr1 dtypetf float32 inputs1tf converttotensorinputarr2 dtypetf float32 self assertallclose inputarr1 outputsoutput0 rtol1e4 atol1e4 self assertallclose inputarr2 outputsoutput1 rtol1e4 atol1e4 def testmultiinputcustommodelandlayerself objectregistration registerkerasserializablepackagemypackage class customlayerlayers layer def callself inputlist self addlossinputlist2 2 return suminputlist objectregistration registerkerasserializablepackagemypackage class custommodelmodels model def buildself inputshape super buildinputshape self layer customlayer tf function def callself inputs inputs listinputs return self layerinputs model custommodel inp tf constanti shape1 1 dtypetf float32 for i in range1 4 expected modelinp path os path joinself gettempdir mykerasmodel tf savedmodel savemodel path restoredmodel tf savedmodel loadpath output restoredmodel callinp self assertallcloseexpected output rtol1e4 atol1e4 def testlisttrackablechildrentrackingself objectregistration registerkerasserializablepackagemypackage class customlayerlistlayers layer def initself super init self sublayers layers dense2 layers dense2 def callself inputs x inputs for sublayer in self sublayers x sublayerx return x inputs layers inputshape1 outputs customlayerlistinputs model models modelinputs outputs inp np array1 0 expected modelinp path os path joinself gettempdir mykerasmodel tf savedmodel savemodel path restoredmodel tf savedmodel loadpath self assertallclose expected restoredmodel signaturesservingdefault tf converttotensorinp dtypetf float32 output0 rtol1e4 atol1e4 def testdicttrackablechildrentrackingself objectregistration registerkerasserializablepackagemypackage class customlayerdictlayers layer def initself super init self sublayers firstlayer layers dense2 secondlayer layers dense2 def callself inputs x inputs for key sublayer in self sublayers items x sublayerx return x inputs layers inputshape1 outputs customlayerdictinputs model models modelinputs outputs inp np array1 0 expected modelinp path os path joinself gettempdir mykerasmodel tf savedmodel savemodel path restoredmodel tf savedmodel loadpath self assertallclose expected restoredmodel signaturesservingdefault tf converttotensorinp dtypetf float32 output0 rtol1e4 atol1e4 def testfixedsignaturestringdtypeself objectregistration registerkerasserializablepackagemypackage class addermodels model tf function inputsignaturetf tensorspecshape dtypetf string def concatself x return x x model adder path os path joinself gettempdir mykerasmodel tf savedmodel savemodel path restoredmodel tf savedmodel loadpath self assertequalmodel concathello restoredmodel concathello def testnonfixedsignaturestringdtypeself objectregistration registerkerasserializablepackagemypackage class addermodels model tf function def concatself x return x x model adder nofnpath os path joinself gettempdir mykerasmodelnofn tf savedmodel savemodel nofnpath restoredmodel tf savedmodel loadnofnpath with self assertraisesregexvalueerror zero restored functions restoredmodel concathello path os path joinself gettempdir mykerasmodel tf savedmodel save model path signaturesmodel concat getconcretefunction tf tensorspecshape dtypetf string namestringinput restoredmodel tf savedmodel loadpath self assertequalmodel concathello restoredmodel concathello def testfinetuningself model customsignaturemodel modelnosignaturespath os path join self gettempdir modelnosignatures modeltf constant0 0 tf savedmodel savemodel modelnosignaturespath restoredmodel tf savedmodel loadmodelnosignaturespath self assertlenlistrestoredmodel signatures keys 0 self assertequalrestoredmodeltf constant3 0 numpy 3 restoredmodel mutatetf constant2 0 self assertequalrestoredmodeltf constant3 0 numpy 6 optimizer optimizers sgd0 05 def trainstep with tf gradienttape as tape loss 10 0 restoredmodeltf constant2 0 2 variables tape watchedvariables grads tape gradientloss variables optimizer applygradientszipgrads variables return loss for in range10 v approaches 5 loss approaches 0 loss trainstep self assertallcloseloss 0 0 rtol1e2 atol1e2 self assertallcloserestoredmodel v numpy 5 0 rtol1e2 atol1e2 def testsignaturespathself model customsignaturemodel modelwithsignaturepath os path join self gettempdir modelwithsignature call model call getconcretefunction tf tensorspecnone tf float32 tf savedmodel savemodel modelwithsignaturepath signaturescall restoredmodel tf savedmodel loadmodelwithsignaturepath self assertequal listrestoredmodel signatures keys servingdefault def testmultiplesignaturesdictpathself model customsignaturemodel modelmultiplesignaturespath os path join self gettempdir modelwithmultiplesignatures call model call getconcretefunction tf tensorspecnone tf float32 signatures servingdefault call arrayinput model call getconcretefunction tf tensorspecnone tf float32 tf savedmodel save model modelmultiplesignaturespath signaturessignatures restoredmodel tf savedmodel loadmodelmultiplesignaturespath self assertequal listrestoredmodel signatures keys servingdefault arrayinput v approaches 5 loss approaches 0
import os

import numpy as np
import pytest
import tensorflow as tf

from keras import backend
from keras import layers
from keras import metrics
from keras import models
from keras import optimizers
from keras import testing
from keras.saving import object_registration


@object_registration.register_keras_serializable(package="my_package")
class CustomModelX(models.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dense1 = layers.Dense(1)
        self.dense2 = layers.Dense(1)

    def call(self, inputs):
        out = self.dense1(inputs)
        return self.dense2(out)

    def one(self):
        return 1


@object_registration.register_keras_serializable(package="my_package")
class CustomSignatureModel(models.Model):
    def __init__(self):
        super(CustomSignatureModel, self).__init__()
        self.v = tf.Variable(1.0)

    @tf.function
    def __call__(self, x):
        return x * self.v

    @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
    def mutate(self, new_v):
        self.v.assign(new_v)


@pytest.mark.skipif(
    backend.backend() != "tensorflow",
    reason="The SavedModel test can only run with TF backend.",
)
class SavedModelTest(testing.TestCase):
    def test_sequential(self):
        model = models.Sequential([layers.Dense(1)])
        model.compile(loss="mse", optimizer="adam")
        X_train = np.random.rand(100, 3)
        y_train = np.random.rand(100, 1)
        model.fit(X_train, y_train)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            model(X_train),
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(X_train, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )

    def test_functional(self):
        inputs = layers.Input(shape=(3,))
        x = layers.Dense(1, name="first_dense")(inputs)
        outputs = layers.Dense(1, name="second_dense")(x)
        model = models.Model(inputs, outputs)
        model.compile(
            optimizer="adam",
            loss="mse",
        )
        X_train = np.random.rand(100, 3)
        y_train = np.random.rand(100, 1)
        model.fit(X_train, y_train)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            model(X_train),
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(X_train, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )

    def test_subclassed(self):
        model = CustomModelX()
        model.compile(
            optimizer="adam",
            loss="mse",
            metrics=[metrics.Hinge(), "mse"],
        )
        X_train = np.random.rand(100, 3)
        y_train = np.random.rand(100, 1)
        model.fit(X_train, y_train)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            model(X_train),
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(X_train, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )

    def test_custom_model_and_layer(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayer(layers.Layer):
            def __call__(self, inputs):
                return inputs

        @object_registration.register_keras_serializable(package="my_package")
        class Model(models.Model):
            def __init__(self):
                super().__init__()
                self.layer = CustomLayer()

            @tf.function(input_signature=[tf.TensorSpec([None, 1])])
            def call(self, inputs):
                return self.layer(inputs)

        model = Model()
        inp = np.array([[1.0]])
        result = model(inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            result,
            restored_model.call(inp),
            rtol=1e-4,
            atol=1e-4,
        )

    def test_multi_input_model(self):
        input_1 = layers.Input(shape=(3,))
        input_2 = layers.Input(shape=(5,))
        model = models.Model([input_1, input_2], [input_1, input_2])
        path = os.path.join(self.get_temp_dir(), "my_keras_model")

        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        input_arr_1 = np.random.random((1, 3)).astype("float32")
        input_arr_2 = np.random.random((1, 5)).astype("float32")

        outputs = restored_model.signatures["serving_default"](
            inputs=tf.convert_to_tensor(input_arr_1, dtype=tf.float32),
            inputs_1=tf.convert_to_tensor(input_arr_2, dtype=tf.float32),
        )

        self.assertAllClose(
            input_arr_1, outputs["output_0"], rtol=1e-4, atol=1e-4
        )
        self.assertAllClose(
            input_arr_2, outputs["output_1"], rtol=1e-4, atol=1e-4
        )

    def test_multi_input_custom_model_and_layer(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayer(layers.Layer):
            def __call__(self, *input_list):
                self.add_loss(input_list[-2] * 2)
                return sum(input_list)

        @object_registration.register_keras_serializable(package="my_package")
        class CustomModel(models.Model):
            def build(self, input_shape):
                super().build(input_shape)
                self.layer = CustomLayer()

            @tf.function
            def call(self, *inputs):
                inputs = list(inputs)
                return self.layer(*inputs)

        model = CustomModel()
        inp = [
            tf.constant(i, shape=[1, 1], dtype=tf.float32) for i in range(1, 4)
        ]
        expected = model(*inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        output = restored_model.call(*inp)
        self.assertAllClose(expected, output, rtol=1e-4, atol=1e-4)

    def test_list_trackable_children_tracking(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayerList(layers.Layer):
            def __init__(self):
                super().__init__()
                self.sublayers = [
                    layers.Dense(2),
                    layers.Dense(2),
                ]

            def call(self, inputs):
                x = inputs
                for sublayer in self.sublayers:
                    x = sublayer(x)
                return x

        inputs = layers.Input(shape=(1,))
        outputs = CustomLayerList()(inputs)
        model = models.Model(inputs, outputs)

        inp = np.array([[1.0]])
        expected = model(inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            expected,
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(inp, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )

    def test_dict_trackable_children_tracking(self):
        @object_registration.register_keras_serializable(package="my_package")
        class CustomLayerDict(layers.Layer):
            def __init__(self):
                super().__init__()
                self.sublayers = {
                    "first_layer": layers.Dense(2),
                    "second_layer": layers.Dense(2),
                }

            def call(self, inputs):
                x = inputs
                for key, sublayer in self.sublayers.items():
                    x = sublayer(x)
                return x

        inputs = layers.Input(shape=(1,))
        outputs = CustomLayerDict()(inputs)
        model = models.Model(inputs, outputs)

        inp = np.array([[1.0]])
        expected = model(inp)
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertAllClose(
            expected,
            restored_model.signatures["serving_default"](
                tf.convert_to_tensor(inp, dtype=tf.float32)
            )["output_0"],
            rtol=1e-4,
            atol=1e-4,
        )

    def test_fixed_signature_string_dtype(self):
        @object_registration.register_keras_serializable(package="my_package")
        class Adder(models.Model):
            @tf.function(
                input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)]
            )
            def concat(self, x):
                return x + x

        model = Adder()
        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(model, path)
        restored_model = tf.saved_model.load(path)
        self.assertEqual(model.concat("hello"), restored_model.concat("hello"))

    def test_non_fixed_signature_string_dtype(self):
        @object_registration.register_keras_serializable(package="my_package")
        class Adder(models.Model):
            @tf.function
            def concat(self, x):
                return x + x

        model = Adder()

        no_fn_path = os.path.join(self.get_temp_dir(), "my_keras_model_no_fn")
        tf.saved_model.save(model, no_fn_path)
        restored_model = tf.saved_model.load(no_fn_path)
        with self.assertRaisesRegex(ValueError, "zero restored functions"):
            _ = restored_model.concat("hello")

        path = os.path.join(self.get_temp_dir(), "my_keras_model")
        tf.saved_model.save(
            model,
            path,
            signatures=model.concat.get_concrete_function(
                tf.TensorSpec(shape=[], dtype=tf.string, name="string_input")
            ),
        )
        restored_model = tf.saved_model.load(path)
        self.assertEqual(model.concat("hello"), restored_model.concat("hello"))

    def test_fine_tuning(self):
        model = CustomSignatureModel()
        model_no_signatures_path = os.path.join(
            self.get_temp_dir(), "model_no_signatures"
        )
        _ = model(tf.constant(0.0))

        tf.saved_model.save(model, model_no_signatures_path)
        restored_model = tf.saved_model.load(model_no_signatures_path)

        self.assertLen(list(restored_model.signatures.keys()), 0)
        self.assertEqual(restored_model(tf.constant(3.0)).numpy(), 3)
        restored_model.mutate(tf.constant(2.0))
        self.assertEqual(restored_model(tf.constant(3.0)).numpy(), 6)
        optimizer = optimizers.SGD(0.05)

        def train_step():
            with tf.GradientTape() as tape:
                loss = (10.0 - restored_model(tf.constant(2.0))) ** 2
            variables = tape.watched_variables()
            grads = tape.gradient(loss, variables)
            optimizer.apply_gradients(zip(grads, variables))
            return loss

        for _ in range(10):
            
            loss = train_step()

        self.assertAllClose(loss, 0.0, rtol=1e-2, atol=1e-2)
        self.assertAllClose(restored_model.v.numpy(), 5.0, rtol=1e-2, atol=1e-2)

    def test_signatures_path(self):
        model = CustomSignatureModel()
        model_with_signature_path = os.path.join(
            self.get_temp_dir(), "model_with_signature"
        )
        call = model.__call__.get_concrete_function(
            tf.TensorSpec(None, tf.float32)
        )

        tf.saved_model.save(model, model_with_signature_path, signatures=call)
        restored_model = tf.saved_model.load(model_with_signature_path)
        self.assertEqual(
            list(restored_model.signatures.keys()), ["serving_default"]
        )

    def test_multiple_signatures_dict_path(self):
        model = CustomSignatureModel()
        model_multiple_signatures_path = os.path.join(
            self.get_temp_dir(), "model_with_multiple_signatures"
        )
        call = model.__call__.get_concrete_function(
            tf.TensorSpec(None, tf.float32)
        )
        signatures = {
            "serving_default": call,
            "array_input": model.__call__.get_concrete_function(
                tf.TensorSpec([None], tf.float32)
            ),
        }

        tf.saved_model.save(
            model, model_multiple_signatures_path, signatures=signatures
        )
        restored_model = tf.saved_model.load(model_multiple_signatures_path)
        self.assertEqual(
            list(restored_model.signatures.keys()),
            ["serving_default", "array_input"],
        )