TensorFlow and Keras 3 Cheatsheet: Models, tf.data and Training

# Install — one wheel covers CPU + GPU on Linux / WSL2 (TF 2.16+)
pip install "tensorflow>=2.17"               # Linux/WSL2 GPU is detected automatically
pip install "tensorflow-cpu>=2.17"           # CPU-only build (smaller)

# macOS — Apple Silicon needs the Metal plugin
pip install tensorflow tensorflow-metal      # MPS-backed accelerator

# Keras 3 as a standalone, multi-backend library (TF / JAX / PyTorch)
pip install "keras>=3"
export KERAS_BACKEND="tensorflow"            # or "jax" / "torch"

# Confirm the device the runtime sees
python -c "
import tensorflow as tf
print('tf', tf.__version__)
print('GPUs:', tf.config.list_physical_devices('GPU'))
print('Keras', tf.keras.__version__)
"

# Memory growth — almost always what you want on a shared GPU
# import tensorflow as tf
# for g in tf.config.list_physical_devices('GPU'):
#     tf.config.experimental.set_memory_growth(g, True)

import tensorflow as tf
import keras
from keras import layers

# 1. Sequential — quickest, for plain feedforward stacks
seq = keras.Sequential([
    layers.Input((28, 28)),
    layers.Flatten(),
    layers.Dense(128, activation="gelu"),
    layers.Dropout(0.2),
    layers.Dense(10),                          # raw logits — let the loss apply softmax
])

# 2. Functional API — branches, residuals, multi-input
inputs  = keras.Input(shape=(28, 28))
x       = layers.Flatten()(inputs)
hidden  = layers.Dense(256, activation="gelu")(x)
skip    = layers.Dense(256, activation="gelu")(hidden) + hidden
outputs = layers.Dense(10)(skip)
func    = keras.Model(inputs, outputs, name="mlp_skip")

# 3. Subclassed Model — full Python control (research code)
class MLP(keras.Model):
    def __init__(self, hidden=128, n_classes=10):
        super().__init__()
        self.flat = layers.Flatten()
        self.fc1  = layers.Dense(hidden, activation="gelu")
        self.drop = layers.Dropout(0.2)
        self.fc2  = layers.Dense(n_classes)

    def call(self, x, training=False):
        x = self.flat(x); x = self.fc1(x)
        x = self.drop(x, training=training)
        return self.fc2(x)

sub = MLP()

# Inspect — works on all three styles after build
sub.build((None, 28, 28))
sub.summary()

import tensorflow as tf

# From in-memory tensors
ds = tf.data.Dataset.from_tensor_slices((X, y))

# From files — variable-length shards, parallel reads
files = tf.data.Dataset.list_files("data/shard-*.tfrec")
ds    = tf.data.TFRecordDataset(files, num_parallel_reads=tf.data.AUTOTUNE)

# Parse one TFRecord example
feature_spec = {
    "image": tf.io.FixedLenFeature([], tf.string),
    "label": tf.io.FixedLenFeature([], tf.int64),
}
def parse(record):
    ex = tf.io.parse_single_example(record, feature_spec)
    img = tf.io.decode_jpeg(ex["image"], channels=3)
    img = tf.image.resize(img, (224, 224)) / 255.0
    return img, ex["label"]

ds = (ds
    .shuffle(10_000, reshuffle_each_iteration=True)
    .map(parse, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(64, drop_remainder=True)
    .prefetch(tf.data.AUTOTUNE)                # overlap CPU prep with GPU train
)

# Caching — for datasets that fit in memory or on disk
ds_cached = ds.cache().repeat()
# ds_cached = ds.cache("/tmp/cache").repeat()  # spill to disk

# Distributed-aware sharding
strategy = tf.distribute.MirroredStrategy()
ds_dist  = strategy.experimental_distribute_dataset(ds)

# End-to-end MNIST — tf.data, Keras model, compile/fit, callbacks, SavedModel.
import tensorflow as tf
import keras
from keras import layers

# Mixed-precision: free speedup on Ampere+ / Apple Silicon
keras.mixed_precision.set_global_policy("mixed_float16")

(x_tr, y_tr), (x_te, y_te) = keras.datasets.mnist.load_data()
x_tr = (x_tr / 255.0).astype("float32"); x_te = (x_te / 255.0).astype("float32")

BATCH = 128
train_ds = (tf.data.Dataset.from_tensor_slices((x_tr, y_tr))
            .shuffle(60_000).batch(BATCH).prefetch(tf.data.AUTOTUNE))
val_ds   = (tf.data.Dataset.from_tensor_slices((x_te, y_te))
            .batch(BATCH).prefetch(tf.data.AUTOTUNE))

model = keras.Sequential([
    layers.Input((28, 28)),
    layers.Flatten(),
    layers.Dense(256, activation="gelu"),
    layers.Dropout(0.2),
    layers.Dense(10, dtype="float32"),         # last layer in fp32 — stable softmax
])

model.compile(
    optimizer=keras.optimizers.AdamW(learning_rate=3e-4, weight_decay=0.01),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["sparse_categorical_accuracy"],
)

cbs = [
    keras.callbacks.EarlyStopping(monitor="val_loss", patience=3, restore_best_weights=True),
    keras.callbacks.ModelCheckpoint("mnist.keras", save_best_only=True, monitor="val_loss"),
    keras.callbacks.TensorBoard(log_dir="./runs"),
]
model.fit(train_ds, validation_data=val_ds, epochs=10, callbacks=cbs)

# Save + reload
model.export("mnist_savedmodel")               # for TF Serving / TFLite
reloaded = keras.models.load_model("mnist.keras")