Keras 3.3.3 Unexpected error when loading the .h5 model

[jeffDebug]

On Windows 10
keras 3.3.3
tensorflow 2.16.1
torch 2.3.0+cu118

I encountered an error when loading the .h5 model, which complained that:

ValueError: A Concatenate layer requires inputs with matching shapes except for the concatenation axis

But the concatenated layers did have the matched sizes and the axis dimension. After the training, the model was saved using the keras.models.Model.save method. This error is consistently reproducible, for the backend of either 'torch' or 'tensorflow'. So I wrote a simple test script.

To provide more context of this issue, we have had many legacy .h5 model files in our application, using Tensorflow 2.x (<= 2.14). We have to load them for the back-compatibility. As Keras 3+ recommended, we may convert and re-save them as .keras files in future.

`
import sys
import keras

import keras.ops as kops
import numpy as np

from keras.layers import BatchNormalization, Concatenate, Conv2D, Dense, Flatten, Input
from keras.models import Model, load_model

TEST_INPUT_CHANNEL = 3
TEST_INPUT_SIZE = 3
TEST_OUTPUT_SIZE = 1
BATCH_SIZE = 32

def make_dummy_model():
inp = Input(
(TEST_INPUT_CHANNEL, TEST_INPUT_SIZE, TEST_INPUT_SIZE),
batch_size=BATCH_SIZE,
name="Input",
)
bn0 = BatchNormalization(axis=1, scale=False, name="bn_0")(inp)

conv0 = Conv2D(
    256,
    (2, 2),
    activation="relu",
    data_format="channels_first",
    name="conv2d_0",
)(bn0)
conv1 = Conv2D(
    128,
    (2, 2),
    activation="relu",
    data_format="channels_first",
    name="conv2d_1",
)(bn0)
concat0 = Concatenate(axis=1, name="concat_0")([conv1, conv0])
bn1 = BatchNormalization(axis=1, scale=False, name="bn_1")(concat0)

conv2 = Conv2D(
    64,
    (2, 2),
    activation="relu",
    data_format="channels_first",
    name="conv2d_2",
)(bn1)
flat0 = Flatten(name="flat_0")(conv2)
oup = Dense(TEST_OUTPUT_SIZE, activation="relu", name="output")(flat0)

model = Model(inputs=inp, outputs=oup)
return model

def f_dummy(x):
return kops.expand_dims(kops.max(x, axis=(1, 2, 3)), axis=1)

if "main" == name:
savepath = sys.argv[1]

# create the keras model
model = make_dummy_model()
model.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.MeanSquaredError(),
)
model.summary()

# train the model
print(f"\nStart training", flush=True)
x = np.random.randn(1024, TEST_INPUT_CHANNEL, TEST_INPUT_SIZE, TEST_INPUT_SIZE)
y = f_dummy(x)
model.fit(x, y, batch_size=BATCH_SIZE, epochs=5)

# save the model
print(f"\nStart saving model {savepath}", flush=True)
model.save(savepath, overwrite=True)
print(f"Finished saving {savepath}", flush=True)

# load the model
try:
    print(f"\nStart loading model {savepath}", flush=True)
    model = load_model(savepath)
    print(f"Done loading model {savepath}", flush=True)
except Exception:
    import traceback

    print(f"\nloading model {savepath} failed.", flush=True)
    print(traceback.format_exc(), flush=True)

`

The following is the full stack trace output.

loading model test.h5 failed.
Traceback (most recent call last):
File "c:\test\test_dummy_model.py", line 85, in
model = load_model(savepath)
File "C:\jeff\conda\envs\keras3\lib\site-packages\keras\src\saving\saving_api.py", line 183, in load_model
return legacy_h5_format.load_model_from_hdf5(
File "C:\jeff\conda\envs\keras3\lib\site-packages\keras\src\legacy\saving\legacy_h5_format.py", line 133, in load_model_from_hdf5
model = saving_utils.model_from_config(
File "C:\jeff\conda\envs\keras3\lib\site-packages\keras\src\legacy\saving\saving_utils.py", line 85, in model_from_config
return serialization.deserialize_keras_object(
File "C:\jeff\conda\envs\keras3\lib\site-packages\keras\src\legacy\saving\serialization.py", line 495, in deserialize_keras_object
deserialized_obj = cls.from_config(
File "C:\jeff\conda\envs\keras3\lib\site-packages\keras\src\models\model.py", line 517, in from_config
return functional_from_config(
File "C:\jeff\conda\envs\keras3\lib\site-packages\keras\src\models\functional.py", line 536, in functional_from_config
process_node(layer, node_data)
File "C:\jeff\conda\envs\keras3\lib\site-packages\keras\src\models\functional.py", line 483, in process_node
layer(*args, **kwargs)
File "C:\jeff\conda\envs\keras3\lib\site-packages\keras\src\utils\traceback_utils.py", line 122, in error_handler
raise e.with_traceback(filtered_tb) from None
File "C:\jeff\conda\envs\keras3\lib\site-packages\keras\src\layers\merging\concatenate.py", line 97, in build
raise ValueError(err_msg)
ValueError: A Concatenate layer requires inputs with matching shapes except for the concatenation axis. Received: input_shape=[(32, 128, 2, 2), (32, 256, 2, 2)]

For the dummy model summary, I'm also attaching the screenshot of the output.

[jeffDebug]

I'm not sure what happened with the python code part, which is segmented in the post. My apologies if I had mis-used the editing tools in the discussion board.

[jeffDebug]

Can you update us with a Google Colab Gist for reproduction? I believe that'd make things easier for us.

Thank You.

Thank you for the feedback. I made a notebook in Colab to demo/reproduce this issue, which is essentially the same code as I posted above. In addition, I found that the issue showed up in keras 3.4.1, too.

Here is the URL of the notebook in my Colab Gist:
https://colab.research.google.com/gist/jeffDebug/149b36c7b9eceb94c8bd9bc773364e2c/untitled0.ipynb

Please let me know if you need any further information from me.

[abhaskumarsinha]

Hello Jeff,

I knew the Concat layer does actually have problems specially when the model is supposed to be saved and loaded.

Have you tried new serialization and deserialization, a doc here?

I usually use a workaround to create a custom layer in Keras and then use concat ops inside it! But I'm not totally sure if that'd work with this case.

Feel free to open a discussion with that Colab gist here.

Thank You.

[jeffDebug]

Hi abhaskumarsinha,

Thank you very much for the response.

The answer to your question about the new saving format ".keras" is, yes I did try that. It worked for both saving and loading the ".keras" format. But I didn't try creating a custom layer.

However neither of the above, if working, would solve our problem. We already had a deployed app that saved client's model in ".h5" format. We have to load any existing ".h5" models, on the client side, for the back-compatibility of the stored models. As I mentioned in the original post, in future we can convert all models to the new ".keras" format. But there is no way we can get around loading the old ".h5" format.

Is there any plan for the near-future release of Keras 3+ to get this bug fixed? If I'm not mistaken, this seems to be a bug that's introduced in the Keras 3+. We had been using Tensorflow up to 2.14.0 (with tensorflow.keras module) saving and loading the model without any issue.

[abhaskumarsinha]

Hello Jeff,

Ah, I understand your concerns. The choice of .h5 format for saving models is rather weird. Technically, .h5 format is supposed to ONLY save the model weights (which is likely some tensors, matrices, compressed using various techniques, sparsity, etc.). You can load an .h5 model ONLY if you have the architecture and the preset available! For example - The model architecture and exact presets define a number of layers and their specifications. If in that case, you have those specifications, then you need a function to define an empty model for you (Let's say - make_test_model()) to return the Keras Model with empty/random initialized weights and then load h5 over them.

In your Colab Notebook, the try-except block would get re-written like this:

# load the model
try:
    print(f"\nStart loading model {savepath}", flush=True)
    model_loaded = make_test_model()
    model_loaded.load_weights(savepath)
    #model_loaded = load_model(savepath)
    print(f"Done loading model {savepath}", flush=True)
    model_loaded.summary()
except Exception:
    import traceback

    print(f"\nloading model {savepath} failed.", flush=True)
    print(traceback.format_exc(), flush=True)

Working Colab Code for Reference

As you can see, I commented your line defining the variable model_loaded using load_model() function that is impossible due to the architectures being not saved in the h5 format to manually myself defining the model using make_test_model() function and then loading h5 weights matrices over it.

Reference: H5- Hierarchical Data Formal - Wikipedia

Best Regards,
Abhas Kumar Sinha

[jeffDebug]

Hi Abhas,

Again thanks so much for the help. That did work and solved the issue.

I was not aware that, the .h5 format saved with the model.save() could be loaded by weights only (model.load_weights()). However, in the previous versions (Tensorflow <= 2.14.0), both the model arch and the weights, in the h5 format, could be loaded by calling keras.models::load_model().

[abhaskumarsinha]

Hello Jeff,

Please note that the same functionality (that was happening with TensorFlow earlier in v2+) is preserved with 3+. It is just merging and concat of the layers where Keras is not able to figure out where to put weights (due to arch complexity) and regular loading of models from .h5 weights could still work in the case of simpler models.

Best Regards

[jeffDebug]

Hi Abhas,

I really appreciate your follow-up.

As I mentioned several time before, the keras.models.::load_model didn't crash in tensorflow 2.x up to 2.14.0 as far as I know. It did work for the Concatenate layers when loading both arch and weights. That's why I suspected it's a bug introduced in the Keras 3+, though I only tried and tested keras 3.3.3 and 3.4.1. I understand that, keras 3+ intended to keep all the same functionality as in tensorflow 2.x.

To show what I said is true, I wrote a similar small notebook in Colab.

Sincerely

[abhaskumarsinha]

Hello Jeff,

Apologies for the oversight. I've raised the issue here: #20019

Please keep a check on the issue. This should get fixed.

Best Regards,
Abhas Kumar Sinha.