Save a Keras Model with a Custom Layer in Python

Recently, while working on a deep learning project in Python, I faced a challenge saving my Keras model that used a custom layer. Everything worked fine during training, but when I tried to reload the model, I got an error saying the custom layer wasn’t recognized.

If you’ve ever built custom layers in Keras and struggled to save or load them properly, you’re not alone. In this tutorial, I’ll show you how to save a Keras model with a custom layer in Python, step by step, with examples you can copy and run right away.

I’ve been working with Keras and TensorFlow for more than four years, and I’ve learned that handling custom layers correctly is crucial when deploying or sharing models.

Table of Contents

Understand the Problem

When you build a deep learning model in Python using Keras, saving and loading it is usually simple:

model.save('model.h5')
new_model = tf.keras.models.load_model('model.h5')

However, this fails when your model includes custom layers, loss functions, or metrics. That’s because Keras doesn’t automatically know how to reconstruct those custom components when reloading the model.

So, we need to tell Keras how to handle our custom objects. There are two main methods to do this: one using the custom_objects argument and another using the get_config() method inside your custom class.

Method 1 – Save and Load Keras Model Using custom_objects

This is the simplest and most direct way to save and load a Keras model with a custom layer in Python.

Step 1: Import Required Python Libraries

We’ll start by importing the necessary packages. I always prefer importing TensorFlow as tf for cleaner code.

import tensorflow as tf
from tensorflow.keras import layers, models

This gives us access to all Keras functionalities, including the base Layer class we’ll use to build our custom layer.

Step 2: Create a Custom Layer in Python

Here’s a simple example of a custom layer that multiplies its input by a constant factor. It’s a minimal but practical demonstration.

class MultiplyLayer(layers.Layer):
    def __init__(self, factor=1.0, **kwargs):
        super(MultiplyLayer, self).__init__(**kwargs)
        self.factor = factor

    def call(self, inputs):
        return inputs * self.factor

This custom layer takes an input tensor and multiplies it by a given factor. You can think of it as a scaling layer.

Step 3: Build a Keras Model with the Custom Layer

Now, let’s integrate this custom layer into a simple sequential model.

input_layer = layers.Input(shape=(1,))
x = MultiplyLayer(factor=3.0)(input_layer)
output_layer = layers.Dense(1)(x)

model = models.Model(inputs=input_layer, outputs=output_layer)
model.compile(optimizer='adam', loss='mse')

This model multiplies the input by 3 and then passes it through a dense layer. It’s simple but demonstrates how custom layers fit into a Keras workflow.

Step 4: Train and Save the Model

We’ll now train the model briefly and then save it. This step ensures the model has some learned parameters.

import numpy as np

# Training data
X = np.array([[1], [2], [3], [4]], dtype=float)
y = np.array([[3], [6], [9], [12]], dtype=float)

model.fit(X, y, epochs=10, verbose=0)

# Save the model
model.save('custom_model.h5')
print("Model saved successfully!")

This saves the model to a file named custom_model.h5. However, if we try to load it directly, we’ll get an error because Keras doesn’t know the MultiplyLayer class yet.

Step 5: Load the Model with custom_objects

To load the model correctly, we must pass our custom layer definition using the custom_objects parameter.

# Load the model
loaded_model = tf.keras.models.load_model(
    'custom_model.h5',
    custom_objects={'MultiplyLayer': MultiplyLayer}
)

print("Model loaded successfully!")

Now, the model loads perfectly because Keras understands the custom layer.

Step 6: Test the Loaded Model

Let’s verify that the loaded model works as expected.

test_input = np.array([[5]], dtype=float)
prediction = loaded_model.predict(test_input)
print("Prediction for input 5:", prediction)

I have executed the above example code and added the screenshot below.

Save Keras Model with Custom Layer Python

If everything is correct, the output should be close to 15, confirming that the custom layer and model were successfully saved and restored.

Method 2 – Save and Load Keras Model Using get_config()

If you plan to share your model or load it without manually specifying custom objects, you can define a get_config() method inside your custom layer. This method tells Keras how to automatically reconstruct the layer.

Step 1: Define a Custom Layer with get_config()

Here’s how I usually define my reusable custom layers in Python.

class MultiplyLayerConfig(layers.Layer):
    def __init__(self, factor=1.0, **kwargs):
        super(MultiplyLayerConfig, self).__init__(**kwargs)
        self.factor = factor

    def call(self, inputs):
        return inputs * self.factor

    def get_config(self):
        config = super(MultiplyLayerConfig, self).get_config()
        config.update({"factor": self.factor})
        return config

This get_config() method ensures that the layer’s parameters (like factor) are stored in the model file.

Step 2: Build, Train, and Save the Model

We can now create and train another model using this improved custom layer.

input_layer = layers.Input(shape=(1,))
x = MultiplyLayerConfig(factor=4.0)(input_layer)
output_layer = layers.Dense(1)(x)

model2 = models.Model(inputs=input_layer, outputs=output_layer)
model2.compile(optimizer='adam', loss='mse')

# Train the model
model2.fit(X, y, epochs=10, verbose=0)

# Save the model
model2.save('custom_model_config.h5')
print("Model with config saved successfully!")

This time, the model saves all configuration details of the custom layer automatically.

Step 3: Load the Model Without custom_objects

Now, since we’ve defined get_config(), Keras can reconstruct the model automatically.

# Load the model
loaded_model2 = tf.keras.models.load_model('custom_model_config.h5', custom_objects={'MultiplyLayerConfig': MultiplyLayerConfig})
print("Model loaded successfully using get_config!")

I have executed the above example code and added the screenshot below.

Save Keras Model with a Custom Layer in Python

Even though we still pass custom_objects, this model can now be shared easily and reloaded by others with minimal setup.

Method 3 – Save Keras Model in a New .keras Format

Keras now supports saving models in the new .keras format, which is more robust and portable than the older .h5 format.

Here’s how you can use it with custom layers.

model.save('custom_model.keras')
loaded_model3 = tf.keras.models.load_model('custom_model.keras', custom_objects={'MultiplyLayerConfig': MultiplyLayerConfig})

The .keras format preserves metadata and configuration better, making it ideal for production environments or TensorFlow Serving.

Common Errors and How to Fix Them

Here are some common issues I’ve faced while saving custom Keras models in Python:

Error: Unknown layer:
This happens when you forget to pass custom_objects during loading. Always include your custom class reference.
Error: AttributeError: ‘MultiplyLayer’ object has no attribute ‘get_config’:
Add a get_config() method to your custom layer class.
Model not loading in different environments:
Make sure the same TensorFlow version is installed, or use the .keras format for better compatibility.

Best Practices When Working with Custom Layers in Python

Always define get_config() in your custom layers for better portability.
Use the .keras format instead of .h5 for modern TensorFlow versions.
Keep your custom layer definitions in a separate Python module for reuse.
Test loading your model immediately after saving to ensure compatibility.
Document your custom layers clearly if you plan to share your model with others.

Saving and loading a Keras model with a custom layer in Python might seem tricky at first, but once you understand the process, it becomes second nature. I’ve used these exact methods in production projects across the USA, especially when deploying models on cloud platforms like AWS and Google Cloud.

The key takeaway is simple: always tell Keras how to rebuild your custom components, either through custom_objects or get_config(). Once you do that, you’ll never face unexpected loading errors again.

You may also like to read the other Keras articles:

Bijay Kumar

I am Bijay Kumar, a Microsoft MVP in SharePoint. Apart from SharePoint, I started working on Python, Machine learning, and artificial intelligence for the last 5 years. During this time I got expertise in various Python libraries also like Tkinter, Pandas, NumPy, Turtle, Django, Matplotlib, Tensorflow, Scipy, Scikit-Learn, etc… for various clients in the United States, Canada, the United Kingdom, Australia, New Zealand, etc. Check out my profile.

enjoysharepoint.com/