Keras 2 to Keras 3 Serialization and Multi-Backend Interoperability

Direct Answer

Loading a Keras 2 model into Keras 3 with a PyTorch or JAX backend does not, by itself, reinitialize weights or introduce numerical drift. If the same architecture and weights are restored correctly, outputs should match the Keras 3 TensorFlow reference within normal floating-point tolerance. Real differences almost always come from a failed restoration, missing custom objects, backend-specific operations, or accumulated float32 kernel differences—not from the act of loading.

Confirmed Facts vs. Likely Causes

Confirmed: the .keras zip format is the portable target for Keras 3. Legacy HDF5 artifacts should be loaded once under TensorFlow, inspected, and re-saved to .keras before cross-backend comparison. Custom layers, losses, and metrics must be supplied through custom_objects. A successful load is not proof of parity.

Likely explanation for observed drift: weights were not actually restored, custom_objects were incomplete, the artifact contained raw TensorFlow ops or Lambda layers that are not backend-agnostic, or backend implementations of convolutions, matrix multiplications, reductions, and batch normalization accumulated float32 rounding differently. Bitwise-identical results across TensorFlow, PyTorch, and JAX are not expected.

Verification Process Without Retraining

1. Freeze the input. Generate one batch with a fixed seed and reuse the exact array in every run.
2. Pin the backend. Set KERAS_BACKEND before importing Keras in each process.
3. Load with custom objects. Supply every legacy custom class or function.
4. Check restoration. Compare parameter count and a hash of concatenated flattened weights before comparing outputs.
5. Compare outputs. Use np.allclose with explicit tolerances and compare argmax decisions separately.
6. Isolate divergence. If mismatch is large, force inference mode, disable dropout and batch-norm updates, and compare layer-by-layer activations to find the first divergent op.
7. Re-save and cross-load. Save the reference model to .keras from one backend and load it in the others; failure there means the artifact is not portable.

import os
os.environ["KERAS_BACKEND"] = "torch"  # or "jax", "tensorflow"
import numpy as np
import keras

# 1. Fixed input
rng = np.random.default_rng(42)
x = rng.normal(size=(8, 224, 224, 3)).astype("float32")
np.save("fixed_input.npy", x)

# 2. Load legacy artifact re-saved to .keras
model = keras.models.load_model(
    "model.keras",
    custom_objects={"MyCustomLayer": MyCustomLayer}
)
model.trainable = False

# 3. Weight integrity check
import hashlib
def weight_hash(m):
    h = hashlib.sha256()
    for w in m.weights:
        h.update(np.asarray(w).astype("float32").tobytes())
    return h.hexdigest()

print("params:", model.count_params())
print("hash:", weight_hash(model))

# 4. Inference and comparison
out = model.predict(x, verbose=0)
ref = np.load("ref_logits.npy")
print("allclose:", np.allclose(ref, out, rtol=1e-5, atol=1e-6))
print("argmax match:", np.array_equal(ref.argmax(1), out.argmax(1)))

One Diagnostic That Changes the Recommendation

Does the Keras 2 artifact or its source code contain custom objects, Lambda layers, or raw TensorFlow ops? If yes, those components must be rewritten as backend-agnostic Keras ops before cross-backend execution is safe; otherwise restrict the model to TensorFlow. If no, the verification steps above are sufficient to confirm parity.

Version Assumptions and Uncertainty

These steps assume Keras 3.x. Legacy load support and op coverage are version-sensitive; confirm exact behavior against your installed minor release. Treat any op-level parity claim as unverified until you run the weight-hash and layer-wise comparison on your own stack.