I’m in the process of migrating to Keras 3. My current version is TensorFlow v2.17.0, and Keras v3.4.1.

I’m testing a bunch of VAE model configurations, which requires me to save these models and load them later. I have tested the models by running the tests on them without saving and loading them. I can verify the models work fine with their weights and configs all accounted for. The problem only happens when the models are loaded from a file. Here is my code for loading a model:

model_name = 'keras3_vae_test'
model_directory = '/content/drive/My Drive/Colab_Files/data/test_models/'
model_path = os.path.join(model_directory, model_name)

# Load the model
model = keras.models.load_model(os.path.join(model_path, 'model_' + model_name + '.keras')

# Load the model history
history_path = os.path.join(model_path, 'history_' + model_name + '.pkl')
if os.path.exists(history_path):
    with open(history_path, 'rb') as file_pi:
        history = pickle.load(file_pi)

print("nModel and history loaded from:", model_path)

The error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-23-2813986d7474> in <cell line: 10>()
      8 
      9 # Load the model
---> 10 model = keras.models.load_model(os.path.join(model_path, 'model_' + model_name + '.keras'), custom_objects=custom_objects)
     11 
     12 # Load the encoder

3 frames
/usr/local/lib/python3.10/dist-packages/keras/src/saving/saving_lib.py in _raise_loading_failure(error_msgs, warn_only)
    454         warnings.warn(msg)
    455     else:
--> 456         raise ValueError(msg)
    457 
    458 

ValueError: A total of 2 objects could not be loaded. Example error message for object <Dense name=z_mean, built=True>:

Layer 'z_mean' expected 2 variables, but received 0 variables during loading. Expected: ['kernel', 'bias']

List of objects that could not be loaded:
[<Dense name=z_mean, built=True>, <Dense name=z_log_var, built=True>]

I believe the issue is due to the serialization of dense layers, but I don’t know how to fix it. Dense layers are common and essential for my model.

Here is a sample of the model I’ve been using for bug testing:

latent_dim = 32

# Encoder
encoder_input = Input(shape=(height, width, channels), name='encoder_input')
x = Conv2D(64, (3, 3), activation='relu', padding='same')(encoder_input)

# Flatten layer
shape_before_flattening = K.int_shape(x)[1:]
x = Flatten()(x)

# Latent space
z_mean = Dense(latent_dim, name='z_mean')(x)
z_log_var = Dense(latent_dim, name='z_log_var')(x)

# Reparameterization trick
@keras.saving.register_keras_serializable()
def sampling(args):
    z_mean, z_log_var = args
    epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim), mean=0., stddev=1.0)
    return z_mean + K.exp(z_log_var / 2) * epsilon

z = Lambda(sampling, output_shape=(latent_dim,), name='z')([z_mean, z_log_var])

# Decoder
decoder_input = Input(K.int_shape(z)[1:])
x = Dense(np.prod(shape_before_flattening))(decoder_input)
x = Reshape(shape_before_flattening)(x)
decoder_output = Conv2D(channels, (3, 3), activation='sigmoid', padding='same')(x)

vae_output = decoder(encoder(encoder_input)[2])
model = Model(encoder_input, vae_output, name='vae')

This is a very basic VAE. After training, I use the following code to save the model:

model_name = 'keras3_vae_test'
model_directory = '/content/drive/My Drive/Colab_Files/data/test_models/'
model_path = os.path.join(model_directory, model_name)
os.makedirs(model_path, exist_ok=True)

# Save the model
model.save(os.path.join(model_path, 'model_' + model_name + '.keras'))

# Save the model history
if 'history' in locals() or 'history' in globals():
    with open(os.path.join(model_path, 'history_' + model_name + '.pkl'), 'wb') as file_pi:
        pickle.dump(history.history, file_pi)

print("nModel and history saved in:", model_path)

As a ‘.keras’ file is essentially just a zip file for the model.config and weights, I tried to save the model in parts:

model_name = 'keras3_vae_test'
model_directory = '/content/drive/My Drive/Colab_Files/data/test_models/'
model_path = os.path.join(model_directory, model_name)

# Load the model architecture and weights
with open(os.path.join(model_path, 'model_' + model_name + '.json'), 'r') as json_file:
    model_json = json_file.read()
model = model_from_json(model_json)
model.load_weights(os.path.join(model_path, 'model_' + model_name + '.weights.h5'))

The error:

/usr/local/lib/python3.10/dist-packages/keras/src/saving/saving_lib.py:576: UserWarning: Skipping variable loading for optimizer 'adam', because it has 34 variables whereas the saved optimizer has 22 variables. 
  saveable.load_own_variables(weights_store.get(inner_path))
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-25-ed9df6bc5f9a> in <cell line: 11>()
      9     model_json = json_file.read()
     10 model = model_from_json(model_json, custom_objects=custom_objects)
---> 11 model.load_weights(os.path.join(model_path, 'model_' + model_name + '.weights.h5'))
     12 
     13 # Load the encoder architecture and weights

1 frames
/usr/local/lib/python3.10/dist-packages/keras/src/saving/saving_lib.py in _raise_loading_failure(error_msgs, warn_only)
    454         warnings.warn(msg)
    455     else:
--> 456         raise ValueError(msg)
    457 
    458 

ValueError: A total of 3 objects could not be loaded. Example error message for object <Conv2D name=conv2d_14, built=True>:

Layer 'conv2d_14' expected 2 variables, but received 0 variables during loading. Expected: ['kernel', 'bias']

List of objects that could not be loaded:
[<Conv2D name=conv2d_14, built=True>, <Dense name=z_mean, built=True>, <Dense name=z_log_var, built=True>]

This gives more details about what’s going on. Specifically, the error is coming from the weights, not all of the optimizer variables are being loaded, and the list of objects that didn’t load. This is more evidence that the dense layers are what Keras 3 is struggling with.