[BUG Report + fix]: Saving after loading tf.keras model

[GaijinOtohp]

Description

It seems that loading the model adds two KeyValuePair to Tensorflow.Train.AutoTrackable. Their keys are non_trainable_variables and layers. Their values are of type Tensorflow.Loader._UserObject.

So saving the model seems to get other KeyValuePair values from TrackableSavedModelSaver.trackable_children(cache) and concatenates them with the KeyValuePair values from Tensorflow.Train.AutoTrackable. This causes a conflict between two values with the same key names (non_trainable_variables and layers) and different values types (Tensorflow.Training.ListWrapper and Tensorflow.Loader._UserObject.)

Converting the concatenated KeyValuePair values to a dictionary will cause the error.

Reproduction Steps

            var layers = tf.keras.layers;

            Tensor input = tf.keras.Input((-1, 1), name: "input_place_holder");
            Tensor hiddenLayer1 = layers.Dense(7, activation: "linear").Apply(input);
            Tensor output = layers.Dense(2, activation: "linear").Apply(hiddenLayer1);
            var model = tf.keras.Model(input, output);

            Tensor results0 = model.predict(tf.constant(new float[,] { { -8 } }));

            model.save("./saved_model/");

            //___________________________________________________________________//
            //___________________________________________________________________//
            model = tf.keras.models.load_model("./saved_model/");

            Tensorflow.NumPy.NDArray inputs = new Tensorflow.NumPy.NDArray(new float[,] { { -5 }, { -4 }, { -3 }, { -2 }, { -1 }, { 0 }, { 1 }, { 2 }, { 3 }, { 4 }, { 5 } });
            Tensorflow.NumPy.NDArray outputs = new Tensorflow.NumPy.NDArray(new float[,] { { 0, 2 }, { 1, 4 }, { 2, 6 }, { 3, 8 }, { 4, 10 }, { 5, 12 }, { 6, 14 }, { 7, 16 }, { 8, 18 }, { 9, 20 }, { 10, 22 } });

            float learningRate = 0.01f;
            model.compile(optimizer: tf.keras.optimizers.SGD(learningRate), loss: tf.keras.losses.MeanSquaredError());

            model.fit(inputs, outputs, epochs: 1000);

            model.save("./saved_model/");

            Tensor results1 = model.predict(tf.constant(new float[,] { { -8 } }));

Known Workarounds

I have just removed the two KeyValuePair values of Tensorflow.Train.AutoTrackable before converting the concatenated KeyValuePair values to dictionary. And it seems to work just fine.

Fix:

In line 30 in the file Tensorflow.Keras/Engine/Layer.Serialize.cs:
return children.Concat(base._trackable_children(save_type, cache)).ToDictionary(x => x.Key, x => x.Value); ---> return children.Concat(base._trackable_children(save_type, cache)).GroupBy(x => x.Key).Select(g => g.First()).ToDictionary(x => x.Key, x => x.Value);

Configuration and Other Information

No response

[Wanglongzhi2001]

Sorry for reply you so late. You can commit a PR to TensorFlow.NET, and your contributions are more than welcome. ^_^

[GaijinOtohp]

That would be my pleasure. Thank you for the suggestion.