What TensorFlow gives you

Three ways to build a model

keras.Sequential([
  Dense(64, activation='relu'),
  Dense(1)
])

x = Dense(64)(inputs)
out = Dense(1)(x)
keras.Model(inputs, out)

class MyModel(keras.Model):
  def call(self, x):
    return self.dense(x)

The Functional API is the recommended default for production code — its explicit, visualisable with keras.utils.plot_model, and serialises cleanly to SavedModel.

import tensorflow as tf
from tensorflow import keras

# Functional API — for multi-input, multi-output, or skip connections
inputs = keras.Input(shape=(128,), name="features")
x = keras.layers.Dense(64, activation="relu")(inputs)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.Dropout(0.3)(x)
x = keras.layers.Dense(32, activation="relu")(x)
outputs = keras.layers.Dense(10, activation="softmax", name="predictions")(x)

model = keras.Model(inputs=inputs, outputs=outputs)
model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)

Keras vs low-level GradientTape

model.fit() handles the training loop, metrics, callbacks, and validation split. Use it unless you need behaviour that Keras cant express. For custom loops — non-standard gradient updates, meta-learning, RL — drop to GradientTape directly.

model.compile(
    optimizer="adam",
    loss="mse",
    metrics=["mae"],
)
model.fit(
    X_train, y_train,
    epochs=20,
    batch_size=64,
    validation_split=0.1,
)

with tf.GradientTape() as tape:
    preds = model(X, training=True)
    loss = loss_fn(y, preds)
grads = tape.gradient(
    loss, model.trainable_variables
)
optimizer.apply_gradients(
    zip(grads, model.trainable_variables)
)

import tensorflow as tf

optimizer = tf.keras.optimizers.Adam(learning_rate=3e-4)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
train_acc = tf.keras.metrics.SparseCategoricalAccuracy()

@tf.function   # compile to a graph for speed
def train_step(X_batch, y_batch):
    with tf.GradientTape() as tape:
        logits = model(X_batch, training=True)
        loss = loss_fn(y_batch, logits)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    train_acc.update_state(y_batch, logits)
    return loss

for epoch in range(num_epochs):
    train_acc.reset_state()
    for X_batch, y_batch in train_dataset:
        loss = train_step(X_batch, y_batch)
    print(f"Epoch {epoch}: acc={train_acc.result():.4f}")

Custom layers

Subclass keras.layers.Layer. Define sub-layers in __init__ and forward logic in call. Pass training=False through to sub-layers that behave differently at inference (BatchNorm, Dropout). Custom layers are fully composable with the Functional API and serialise correctly.

import tensorflow as tf
from tensorflow import keras

class ResidualBlock(keras.layers.Layer):
    def __init__(self, units: int, **kwargs):
        super().__init__(**kwargs)
        self.dense1 = keras.layers.Dense(units, activation="relu")
        self.dense2 = keras.layers.Dense(units)
        self.norm = keras.layers.LayerNormalization()
        self.add = keras.layers.Add()

    def call(self, x, training=False):
        residual = x
        x = self.dense1(x)
        x = self.dense2(x)
        x = self.add([x, residual])
        return self.norm(x, training=training)

# Use inside any model
inputs = keras.Input(shape=(64,))
x = ResidualBlock(64)(inputs)
x = ResidualBlock(64)(x)
outputs = keras.layers.Dense(1)(x)
model = keras.Model(inputs, outputs)

Data pipelines with tf.data

tf.data is TensorFlow s declarative data pipeline API. It chains transformations lazily and can overlap CPU preprocessing with GPU compute via prefetch(AUTOTUNE). Always prefer it over feeding NumPy arrays for anything beyond small toy datasets.

import tensorflow as tf

AUTOTUNE = tf.data.AUTOTUNE

def parse_example(path, label):
    img = tf.io.read_file(path)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, [224, 224])
    img = tf.cast(img, tf.float32) / 255.0
    return img, label

train_ds = (
    tf.data.Dataset.from_tensor_slices((image_paths, labels))
    .shuffle(buffer_size=1000)
    .map(parse_example, num_parallel_calls=AUTOTUNE)
    .batch(32)
    .prefetch(AUTOTUNE)   # overlap preprocessing + GPU compute
)

Callbacks

Callbacks hook into model.fit() at epoch start/end, batch start/end, and on training end. Always use at least EarlyStopping + ModelCheckpoint in production training runs.

import tensorflow as tf
from tensorflow import keras

callbacks = [
    # Stop early and restore best weights
    keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=5, restore_best_weights=True
    ),
    # Save best checkpoint to disk
    keras.callbacks.ModelCheckpoint(
        "best_model.keras", monitor="val_loss", save_best_only=True
    ),
    # Reduce LR when plateau
    keras.callbacks.ReduceLROnPlateau(
        monitor="val_loss", factor=0.5, patience=3, min_lr=1e-6
    ),
    # TensorBoard logging
    keras.callbacks.TensorBoard(log_dir="./logs", histogram_freq=1),
]

model.fit(X_train, y_train, epochs=100, validation_split=0.1, callbacks=callbacks)

For custom behaviour, subclass keras.callbacks.Callback and override on_epoch_end, on_batch_end, etc. Useful for logging to experiment trackers (W&B, MLflow) or early stopping on custom metrics.

Transfer learning & fine-tuning

keras.applications ships pretrained image models (MobileNet, EfficientNet, ResNet, etc.) with ImageNet weights. Freeze the base, train a new head, then unfreeze the top layers for fine-tuning at a very low learning rate.

import tensorflow as tf
from tensorflow import keras

# Load pretrained base (ImageNet weights, no top classifier)
base = keras.applications.MobileNetV3Small(
    input_shape=(224, 224, 3),
    include_top=False,
    weights="imagenet",
)
base.trainable = False   # freeze base during initial training

# Add custom head
inputs = keras.Input(shape=(224, 224, 3))
x = keras.applications.mobilenet_v3.preprocess_input(inputs)
x = base(x, training=False)
x = keras.layers.GlobalAveragePooling2D()(x)
x = keras.layers.Dropout(0.2)(x)
outputs = keras.layers.Dense(num_classes, activation="softmax")(x)
model = keras.Model(inputs, outputs)

# Phase 1: train head only
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])
model.fit(train_ds, epochs=10, validation_data=val_ds)

# Phase 2: unfreeze top layers and fine-tune
base.trainable = True
for layer in base.layers[:-20]:   # keep early layers frozen
    layer.trainable = False

model.compile(
    optimizer=keras.optimizers.Adam(1e-5),   # very low LR for fine-tuning
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)
model.fit(train_ds, epochs=10, validation_data=val_ds)

Saving, loading & exporting

The .keras format (Keras native) is the recommended save format — it stores architecture, weights, and optimizer state. Use SavedModel for TF Serving. Use TFLite for mobile and edge. Always include preprocessing layers inside the model before exporting so the serving artifact is self-contained.

# Save the full model (architecture + weights + optimizer state)
model.save("my_model.keras")             # Keras native format (recommended)
model.save("my_model_savedmodel")        # TF SavedModel format (for TF Serving)

# Load
loaded = tf.keras.models.load_model("my_model.keras")

# Export for mobile / edge — TensorFlow Lite
converter = tf.lite.TFLiteConverter.from_saved_model("my_model_savedmodel")
converter.optimizations = [tf.lite.Optimize.DEFAULT]   # quantization
tflite_model = converter.convert()
with open("model.tflite", "wb") as f:
    f.write(tflite_model)

# Weights only (transfer learning)
model.save_weights("weights.h5")
model.load_weights("weights.h5")

Common layers reference