Bug: Exported/saved models are completely inaccurate for inference

[g-rov]

Bug Description

Models exported with export_model() and exported models that are saved and loaded are completely inaccurate for inference. This affects StructuredDataRegressor and most probably other models that rely on pipelines.
In my understanding, this is because the export_model() function does not take into account any pre- and post-processors that might be included in the pipeline. Hence, the predictions are off.

Bug Reproduction

Code for reproducing the bug:

import tensorflow as tf
import keras
import numpy as np
import autokeras as ak
    
# Example usage
if __name__ == "__main__":
    max_trials = 10
    epochs = 20
    batch_size = 50

    X = np.random.rand(200,10)
    X[:,:5] = X[:,:5] * [1, 5, 100, 2, 0.1] + [0, 2, 1000, -200, -1]
    # Creating categorical variables to force encoder
    X[:,6] = np.round(X[:,6])
    X[:,7:] = np.round(X[:,7:],1)
    X_train = X[:100]
    X_test = X[100:150]
    X_val = X[150:]

    y = np.random.rand(200,1) * 20 + 50
    y_train = y[:100]
    y_val = y[100:150]
    y_test = y[150:]

    model = ak.StructuredDataRegressor(
        column_names=["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"],
        column_types={"0": "numerical", "1": "numerical", "2": "numerical", "3": "numerical", "4": "numerical", "5": "categorical", "6": "categorical", "7": "categorical", "8": "categorical", "9": "categorical"},
        output_dim=1,
        loss="mse",
        metrics=[keras.metrics.RootMeanSquaredError()],
        project_name="structured_data_regressor",
        max_trials=max_trials,
        directory='structured_data_regressor',
        objective="val_loss",
        overwrite=False,
        seed=42,
    )
    model.fit(
        X_train,
        y_train,
        validation_data=(X_val, y_val),
        epochs=epochs,
        batch_size=batch_size,
        verbose=2
    )
    best_model = model.export_model()
    best_model.save('best_model.keras')

    # Evaluate AutoModel
    val_loss, val_rmse = model.evaluate(X_test, y_test, batch_size=batch_size, verbose=2)
    print(f"\nValidation Weighted MSE: {val_loss:.6f}")
    print(f"Validation RMSE: {val_rmse:.6f}")

    # Evaluate best model
    model_backup = model
    del model
    model = best_model
    val_loss, val_rmse = model.evaluate(X_test, y_test, batch_size=batch_size, verbose=2)
    print(f"\nValidation Weighted MSE: {val_loss:.6f}")
    print(f"Validation RMSE: {val_rmse:.6f}")

    # Evaluate best reloaded model
    del model
    model = keras.models.load_model('best_model.keras')
    val_loss, val_rmse = model.evaluate(X_test, y_test, batch_size=batch_size, verbose=0)
    print(f"\nValidation Weighted MSE: {val_loss:.6f}")
    print(f"Validation RMSE: {val_rmse:.6f}")

Expected Behavior

The MSE / RMSE reported by the StructuredDataRegressor, the exported model and the reloaded model should match. However, they differ significantly. In the example code above, the difference is consistently ~20%. Depending on the data (number of categorical variables, distribution, etc.) the difference can go up by orders of magnitude.

Setup Details

Include the details about the versions of:

• OS type and version: Ubuntu 24.04.3 LTS
• Python: 3.13.9
• autokeras: 3.0.0
• keras-tuner: 1.4.8
• scikit-learn: 1.7.2
• numpy: 2.3.5
• pandas: 2.3.3
• tensorflow: 2.20.0

Additional context

The export_model() function should produce a model that has equal or, at least, equivalent performance with the AutoModel.

[certainly-cyber]

Yes, I noticed the same issue as well. Here is my reproduction code:

# Set Keras backend to torch BEFORE importing any keras/autokeras modules
os.environ["KERAS_BACKEND"] = "torch"

import pandas as pd
import numpy as np
import autokeras as ak
import keras
import torch

# 1. Data Preparation
df = pd.read_csv('california_housing.csv')
x = df.drop(columns=['target']).values
y = df['target'].values

train_size = int(len(x) * 0.9)
x_train, y_train = x[:train_size], y[:train_size]
x_test, y_test = x[train_size:], y[train_size:]

# 2. Initialize and Train
# Note: Max_trials is set to 3 for quick demo
reg = ak.StructuredDataRegressor(overwrite=True, max_trials=3)
reg.fit(x_train, y_train, epochs=1)

# 3. Predict and Evaluate (Original AutoKeras Object)
original_predicted_y = reg.predict(x_test)
original_eval_results = reg.evaluate(x_test, y_test)
original_loss = original_eval_results[0]
original_mse = original_eval_results[1]

# 4. Save the Model
# Export the model as a Keras Model (running on torch backend)
model = reg.export_model()
model_path = "best_housing_model_torch.keras"

try:
    model.save(model_path)
    print(f"\n[OK] Model saved successfully with PyTorch backend: {model_path}")
except Exception as e:
    print(f"\n[Error] Failed to save model: {e}")

# 5. Load the Model
# Under Keras 3, this loads the model into the current torch backend
loaded_model = keras.models.load_model(
    model_path, 
    custom_objects=ak.CUSTOM_OBJECTS
)

# 6. Verify and Compare (PyTorch Consistency)
# In PyTorch backend, we can check if it's an instance of torch.nn.Module
print(f"Is loaded model a PyTorch module? {isinstance(loaded_model, torch.nn.Module)}")

loaded_eval_results = loaded_model.evaluate(x_test, y_test, verbose=0)
loaded_loss = loaded_eval_results[0]
loaded_mse = loaded_eval_results[1]

loaded_predicted_y = loaded_model.predict(x_test)

print("\n" + "="*60)
print("             PyTorch Backend Consistency Verification")
print("="*60)

print(f"{'Metric':<10} | {'Original (AutoKeras)':<20} | {'Loaded (PyTorch)':<15} | {'Difference':<10}")
print("-" * 60)
print(f"{'Loss':<10} | {original_loss:<20.6f} | {loaded_loss:<15.6f} | {abs(original_loss - loaded_loss):.2e}")
print(f"{'MSE':<10} | {original_mse:<20.6f} | {loaded_mse:<15.6f} | {abs(original_mse - loaded_mse):.2e}")

# Consistency Checks
loss_consistent = np.isclose(original_loss, loaded_loss, atol=1e-5)
predict_consistent = np.allclose(original_predicted_y.flatten(), loaded_predicted_y.flatten(), atol=1e-5)

print(f"\nIs Test Set Loss consistent? {loss_consistent}")
print(f"Are prediction results consistent? {predict_consistent}") 

I have identified a significant consistency issue when exporting and reloading models from AutoKeras using Keras 3. This issue persists regardless of the backend used (both PyTorch and TensorFlow). When the model is reloaded from a .keras file, the evaluation metrics and predictions differ drastically from the original trained model. In many cases, the loaded model behaves as if it has not been trained at all.

Furthermore, I found that if I manually construct the search space and strictly prohibit the use of any layers with moving average states (such as BatchNormalization), the results become consistent. Therefore, it is likely an issue with how the state of these specific layers is saved.