[BUG Report]: Normalization layer, must call .adapt after load and save

[mvphelps]

Description

I believe a normalization layer is not being loaded correctly when loading a saved model from disk.

I have a multiclass classification problem. My model doesn't work after saving and loading. I've created code to reproduce the issue, included below. I've found the problem has to do with including a normalization layer in my sequential model. The attached code demonstrates the issue. Note that it outputs the weights and the weights appear to be loaded correctly. However the model when loaded from disk only shows 0.3 accuracy (it should be 1). If I then call .adapt again on the normalization layer, with the original training data, evaluating the model again works properly.

The repro code uses a trivial dataset with 3 1 hot encoded features. To clarify the output - the model is evaluated 3 times. The first time it shows 1 for accuracy as it is fitting perfectly. After saving and reloading the model, the reloaded model is evaluated and show 0.3 accuracy. Lastly, the 3rd evaluation shows matching the original, after a hack to invoke .adapt on the normalization layer again.

See the TODOs in the code. You can remove the normalization layer, and the loaded model accuracy will match the original exactly. Also I have it saving the model that was loaded with a different name. There are some differences in the meta data and it appears to be adding extra normalization layers when loading.

Please note I am aware that this toy dataset does not need a normalization layer. However my actual use case does. This code is just highly simplified to mimic my scenario as closely as I can for you.

Reproduction Steps

New console program, add this to main.

var m = new MulticlassTest();
m.Run();
Console.WriteLine();
Console.WriteLine();
Console.WriteLine("Done, Press <Enter> to quit");
Console.ReadLine();

Add this class program.

using Tensorflow.Keras;
using static Tensorflow.Binding;
using static Tensorflow.KerasApi;
using Tensorflow.Keras.Engine;
using Tensorflow.NumPy;
using Tensorflow.Keras.Layers;
using Tensorflow;
using System.Text;

namespace tfnet_repro;

public class MulticlassTest
{
    NDArray x_train, y_train, x_test, y_test;

    public bool Run()
    {
        tf.set_random_seed(10);
        tf.enable_eager_execution();

        PrepareData();

        var (originalModel, originalWeights) = Train(BuildModel(3, 3));
        var (loadedModel, loadedWeights) = LoadAndReEvaluate();
        var (adaptedModel, adaptedWeights) = ReAdaptAndReEvaluate(loadedModel);

        Console.WriteLine($"\r\n Weights");
        Console.WriteLine($"Original:  {originalWeights}");
        Console.WriteLine($"Loaded:    {loadedWeights}");
        Console.WriteLine($"Readapted: {adaptedWeights}");

        if (originalWeights == loadedWeights)
        {
            Console.WriteLine("\r\nLayers and Weights are same");
        }

        var originalModelText = DebugLayers(originalModel);
        var loadedModelText = DebugLayers(loadedModel);
        if (originalModelText == loadedModelText)
        {
            Console.WriteLine("\r\nOriginal and Loaded models appear identical");
        }
        return true;
    }
    private record TrainingEntry(float Label, float[] Features);

    public void PrepareData()
    {   //Data is all one-hot. Label corresponds to the position of the 1.
        int j = 0;
        var entries = new List<TrainingEntry>();
        for (int i = 0; i < 100; i++) entries.Add(new TrainingEntry(0, new float[] { 1, 0, 0 }));
        for (int i = 0; i < 100; i++) entries.Add(new TrainingEntry(1, new float[] { 0, 1, 0 }));
        for (int i = 0; i < 100; i++) entries.Add(new TrainingEntry(2, new float[] { 0, 0, 1 }));
        entries.Shuffle(3);
        var features = entries.Select(x => x.Features).ToList().To2DArray();
        var labels = entries.Select(x=>x.Label).ToArray();
        (x_train, y_train, x_test, y_test) = train_test_split(features, labels, 0.1f);
    }
    public static (NDArray, NDArray, NDArray, NDArray) train_test_split(NDArray x, NDArray y, float test_size = 0.2f)
    {
        var len = x.shape[0];
        int train_size = (int)Math.Round(len * (1 - test_size));
        var train_x = x[new Slice(stop: train_size), new Slice()];
        var test_x = x[new Slice(start: train_size), new Slice()];
        var train_y = y[new Slice(stop: train_size)];
        var test_y = y[new Slice(start: train_size)];

        return (train_x, train_y, test_x, test_y);
    }

    public IModel BuildModel(int numberOfLabels, int numberOfFeatures)
    {
        var layers = new LayersApi();
        var normalizer = layers.Normalization(x_train.shape).Named("CustomNorm");
        normalizer.adapt(x_train);

        var model = keras.Sequential(new List<ILayer>()
        {
            // TODO: Save and reload works fine if you comment out the normalizer and just
            // bypass it (but the model doesn't train well)
            normalizer,
            layers.Dense(numberOfFeatures, activation: KerasApi.keras.activations.Relu),
            layers.Dense(numberOfLabels).Named("Output")
        });
        model.summary();
        CompileModel(model);
        return model;
    }

    public (IModel model, string weights) Train(IModel model)
    {
        model.fit(x_train, y_train, batch_size: 64, epochs: 31, validation_split: 0.1f, verbose: 0);
        Console.WriteLine("\r\nOriginal Built Model");
        model.evaluate(x_test, y_test, verbose: 2);
        model.save(nameof(MulticlassTest));
        var weights = DebugLayerWeights(model.Layers[0]);
        return (model, weights);
    }

    public (IModel newModel, string loadedWeights) LoadAndReEvaluate()
    {
        Console.WriteLine("\r\n\r\nLoading from disk");
        var newModel = keras.models.load_model(nameof(MulticlassTest), true);
        CompileModel(newModel); //If we don't recompile the model here, it throws a NullReferenceException om .evaluate later
        var loadedWeights = DebugLayerWeights(newModel.Layers[0]);
        Console.WriteLine("Reloaded Model ------------ Accuracy is wrong!!!");
        newModel.evaluate(x_test, y_test, verbose: 2);
        newModel.summary();

        // TODO: Save a copy of the loaded model, and compare the keras_metadata.pb files - the second
        // saved model has an extra normalization layer in it in both the input_layers and output_layers
        // sections. These models should be identical since we've made no changes.
        newModel.save(nameof(MulticlassTest) + "2");

        return (newModel, loadedWeights);
    }

    public (IModel model, string readaptedWeights) ReAdaptAndReEvaluate(IModel model)
    {
        var normalizer = model.Layers[0];
        normalizer.adapt(x_train);
        Console.WriteLine("\r\n\r\nRe-Adapted normalization layer with reloaded model");
        model.evaluate(x_test, y_test, verbose: 2);
        var readaptedWeights = DebugLayerWeights(normalizer);
        return (model, readaptedWeights);
    }
    private void CompileModel(IModel newModel)
    {
        newModel.compile(loss: keras.losses.SparseCategoricalCrossentropy(from_logits: true),
            optimizer: keras.optimizers.Adam(),
            metrics: new[] { "accuracy" });
    }
    private string DebugLayers(IModel model)
    {
        var sb = new StringBuilder();
        foreach (var layer in model.Layers)
        {
            sb.AppendLine(DebugLayerWeights(layer));
        }

        return sb.ToString();
    }

    private string DebugLayerWeights(ILayer layer)
    {
        return $"{layer.Name} = {String.Join(',', layer.get_weights())}";
    }
}
public static class Extensions
{
    public static T[,] To2DArray<T>(this IList<T[]> source)
    {   //From Jon Skeet, https://stackoverflow.com/questions/9774901/how-to-convert-list-of-arrays-into-a-multidimensional-array
        if (source.Count == 0)
        {
            return new T[,] { };
        }
        int minorLength = source[0].Length;
        if (minorLength == 0)
        {
            return new T[,] { };
        }
        T[,] ret = new T[source.Count, minorLength];
        for (int i = 0; i < source.Count; i++)
        {
            var array = source[i];
            if (array.Length != minorLength)
            {
                throw new ArgumentException("All minor arrays must be the same length");
            }
            for (int j = 0; j < minorLength; j++)
            {
                ret[i, j] = array[j];
            }
        }
        return ret;
    }
    public static void Shuffle<T>(this IList<T> list, int randomSeed = 0)
    {   //Adapted to be deterministic (accepts seed), from https://stackoverflow.com/questions/273313/randomize-a-listt
        int n = list.Count;
        var rand = new Random(randomSeed);
        while (n > 1)
        {
            n--;
            int k = rand.Next(n + 1);
            (list[k], list[n]) = (list[n], list[k]);
        }
    }
    public static ILayer Named(this ILayer source, string name)
    {
        ((Layer)source).Name = name;
        return source;
    }
}

Known Workarounds

Calling adapt on the layer again after reloading works, but this is effectively training the layer again, and would require shipping training data to prod.

Is there a step for saving and loading I am missing?

Configuration and Other Information

Latest versions of Tensorflow.Net, and version 2.16 of Redist. Running on Windows 10 with dotnet 8.