Practical: Same Model on All Three Backends

The whole track, in one script

This is where the promise stops being a claim and becomes something you can run. The model definition, the compile/fit loop, and the save call in the Code block are 100% backend-agnostic — not one line names TensorFlow, PyTorch, or JAX. You choose the engine with the KERAS_BACKEND env var outside this file, and the same source trains to roughly the same accuracy on all three.

Why .keras is the linchpin

The final model.save("my_model.keras") is the payoff. The .keras format stores architecture, weights, and optimizer state with no backend-specific operations, so you can train under JAX for speed and reload under TensorFlow for deployment without retraining or converting anything. Save once, load anywhere.

What will differ between the three runs is the stuff worth measuring yourself: per-epoch wall time and peak memory. That's exactly what the exercise asks you to record — the only benchmark that should drive your backend choice is your model on your hardware.

Code

One backend-agnostic script: build, train, save·python

import keras
from keras import layers

# This model definition is 100% backend-agnostic
def build_model():
    inputs = keras.Input(shape=(784,))
    x = layers.Dense(256, activation="relu")(inputs)
    x = layers.Dropout(0.3)(x)
    x = layers.Dense(128, activation="relu")(x)
    x = layers.Dropout(0.3)(x)
    outputs = layers.Dense(10, activation="softmax")(x)
    return keras.Model(inputs, outputs)

model = build_model()
model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)

# Train on any backend — same result
model.fit(x_train, y_train, epochs=10, validation_split=0.2)

# Save — cross-backend compatible!
model.save("my_model.keras")
# Load this on ANY backend — TF, PyTorch, JAX

Practical: Same Model on All Three Backends

The whole track, in one script

Why .keras is the linchpin

Code

External links

Exercise

Progress

Comments 0