How to create a subclassed model in Keras

Subclassed Keras models are useful when a network needs Python control flow or a custom forward path that does not fit a plain Sequential stack or Functional API graph. A subclass keeps layer objects on the model and puts the forward pass in call(), while still using the built-in compile(), fit(), predict(), and trainable-variable tracking APIs.

The class should create child layers in __init__() and use them in call(). Keras tracks layers assigned as attributes, and the first model call builds their variables for the observed input shape. Keep backend-specific operations out of call() unless the project is intentionally tied to one backend.

A small scoring model on the JAX backend exercises the subclass through model building, one training epoch, and prediction. Replace the sample arrays and layer sizes with project data after the class shape, variable tracking, and output shape are proven.

1. Create subclass_model.py with backend selection, imports, and sample data.

   subclass_model.py

   import os
    
   os.environ["KERAS_BACKEND"] = "jax"
    
   import keras
   import numpy as np
    
   keras.utils.set_random_seed(11)
    
   x_train = np.array(
       [
           [0.0, 0.0, 0.1],
           [0.0, 1.0, 0.4],
           [1.0, 0.0, 0.6],
           [1.0, 1.0, 0.9],
       ],
       dtype="float32",
   )
   y_train = np.array([0.0, 1.0, 1.0, 0.0], dtype="float32")

   Set KERAS_BACKEND before the first import keras. With KERAS_BACKEND set to jax, the Python environment must have both Keras and JAX installed.
   Related: How to install Keras with pip
   Related: How to set the Keras backend

2. Add a Model subclass that creates layers in __init__() and applies them in call().

   class FraudScoreModel(keras.Model):
       def __init__(self, hidden_units=8):
           super().__init__(name="fraud_score_model")
           self.hidden = keras.layers.Dense(hidden_units, activation="relu")
           self.dropout = keras.layers.Dropout(0.1)
           self.score = keras.layers.Dense(1, activation="sigmoid")
    
       def call(self, inputs, training=False):
           x = self.hidden(inputs)
           x = self.dropout(x, training=training)
           return self.score(x)

   The training argument lets layers such as Dropout behave differently during fit() and prediction without changing the public model call.

3. Instantiate the subclass and call it on a sample batch.

   model = FraudScoreModel(hidden_units=8)
    
   sample_scores = model(x_train[:2], training=False)

   The first call builds the child layer variables. For subclassed models, call the model once with representative input before inspecting build state or trainable variables.

4. Compile the subclassed model for a binary target.

   model.compile(
       optimizer=keras.optimizers.Adam(learning_rate=0.05),
       loss=keras.losses.BinaryCrossentropy(),
       metrics=[keras.metrics.BinaryAccuracy(name="accuracy")],
   )

   A subclassed model can use compile() when call() returns a tensor that matches the selected loss and target shape.
   Related: How to compile a model in Keras

5. Train the model for one short epoch.

   history = model.fit(x_train, y_train, epochs=1, batch_size=4, verbose=0)

6. Run prediction on a small batch after training.

   predictions = model.predict(x_train[:2], verbose=0)

7. Print the build, variable, training, and prediction evidence.

   print(f"backend: {keras.backend.backend()}")
   print(f"model name: {model.name}")
   print(f"built: {model.built}")
   print(f"sample output shape: {sample_scores.shape}")
   print(f"trainable variables: {len(model.trainable_variables)}")
   print("history keys:", ", ".join(sorted(history.history.keys())))
   print(f"prediction shape: {predictions.shape}")
   print("prediction sample:", [float(f"{value:.4f}") for value in predictions[:, 0]])

8. Run the script and confirm that the subclass builds, trains, and predicts.

   $ python subclass_model.py
   backend: jax
   model name: fraud_score_model
   built: True
   sample output shape: (2, 1)
   trainable variables: 4
   history keys: accuracy, loss
   prediction shape: (2, 1)
   prediction sample: [0.4741, 0.5877]