Practical: Save → Cross-Backend Load → Export

The whole track in one run

This lesson stitches together everything the track introduced separately. The .keras format (lesson 1), cross-backend loading (lesson 4), and export (lesson 5) aren't three unrelated tricks — they're three stages of one pipeline that carries a model from a fast training engine to a stable serving runtime without ever retraining it.

Three stages, three contexts

Stage 1 — train where it's fastest. Set the JAX backend, build and fit the model, and write a single model.keras. The optimizer state rides along, so this checkpoint is fully resumable. Stage 2 — load where it's stable. In a separate process with the TensorFlow backend selected, load_model reads the same file; the backend-agnostic weights mean nothing is lost in the handoff. Stage 3 — export for production. Project the loaded model onto a serving format with export, and ship it through TF Serving, TFLite, or TF.js.

Why the backend switch happens between processes

Notice each stage is its own script with its own KERAS_BACKEND. The backend is locked once import keras runs, so you can't flip engines mid-program — the switch lives at the process boundary, and the .keras file is what carries the model across that boundary intact. That file is the contract; the backends on either side are interchangeable.

Code

Train (JAX) → save → load (TF) → export for production·python

# Step 1: Train on JAX
os.environ["KERAS_BACKEND"] = "jax"
import keras
model = build_model()
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
model.fit(x_train, y_train, epochs=10)
model.save("model.keras")

# Step 2: Load on TensorFlow
# (In a new script with KERAS_BACKEND="tensorflow")
model = keras.models.load_model("model.keras")

# Step 3: Export for production
model.export("production_model", format="tf_saved_model")
# Deploy with TF Serving, TFLite, or TF.js

Practical: Save → Cross-Backend Load → Export

The whole track in one run

Three stages, three contexts

Why the backend switch happens between processes

Code

External links

Exercise

Progress

Comments 0