TensorFlow Cheatsheet

Training (compile, fit)

Use this TensorFlow reference while you build software engineering projects, review code for technical interview prep, or polish examples for a software engineer resume.

model.compile()

model.compile(
    optimizer='adam',                          # string, instance, or tf.optimizers.*
    loss='sparse_categorical_crossentropy',    # string, function, or Loss instance
    metrics=['accuracy', 'AUC'],               # evaluated each epoch
    loss_weights=None,                         # dict or list for multi-output
    weighted_metrics=None,                     # metrics applied to sample_weight
    run_eagerly=False,                         # True to debug (disables tf.function)
    steps_per_execution=1,                     # batch calls per graph step (TPU perf)
    jit_compile=False,                         # XLA compilation
)

# Multi-output example
model.compile(
    optimizer=keras.optimizers.Adam(1e-3),
    loss={'head_a': 'mse', 'head_b': 'binary_crossentropy'},
    loss_weights={'head_a': 1.0, 'head_b': 2.0},
    metrics={'head_b': ['accuracy', 'AUC']},
)

model.fit()

history = model.fit(
    x=train_dataset,           # numpy, tf.Tensor, tf.data.Dataset, dict, generator
    y=None,                    # labels (not needed if x is a Dataset)
    batch_size=32,             # ignored when x is a Dataset
    epochs=10,
    verbose='auto',            # 0=silent, 1=progress bar, 2=one line per epoch
    callbacks=None,            # list of keras.callbacks.*
    validation_split=0.1,      # fraction from training data (not for Dataset)
    validation_data=(x_val, y_val),   # or a Dataset
    shuffle=True,
    class_weight=None,         # dict {class_id: weight}
    sample_weight=None,        # per-sample weights
    initial_epoch=0,
    steps_per_epoch=None,      # inferred from Dataset len if None
    validation_steps=None,
    validation_batch_size=None,
    validation_freq=1,         # int or list of epochs to run val on
)

History Object

history.history           # dict: {'loss': [...], 'val_loss': [...], 'accuracy': [...]}
history.epoch             # list of epoch indices
history.params            # dict: {'epochs': 10, 'steps': ..., 'verbose': ...}

import matplotlib.pyplot as plt
plt.plot(history.history['loss'], label='train')
plt.plot(history.history['val_loss'], label='val')

model.evaluate()

loss, *metrics = model.evaluate(
    x_test, y_test,
    batch_size=32,
    verbose=1,
    return_dict=False,     # True → returns dict instead of list
)

results = model.evaluate(x_test, y_test, return_dict=True)
# {'loss': 0.42, 'accuracy': 0.91}

Custom Training Loop

For full control over gradient computation, update order, or multi-model training.

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

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss   = loss_fn(y, logits)
        loss   = loss + sum(model.losses)   # regularization terms
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    train_acc.update_state(y, logits)
    return loss

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

Validation Step

val_acc = keras.metrics.SparseCategoricalAccuracy()

@tf.function
def val_step(x, y):
    logits = model(x, training=False)
    val_acc.update_state(y, logits)

for epoch in range(epochs):
    ...
    val_acc.reset_state()
    for x_v, y_v in val_dataset:
        val_step(x_v, y_v)
    print(f'Val acc: {val_acc.result():.4f}')

Custom train_step (subclass override)

Override train_step to keep model.fit() conveniences (callbacks, history) while customizing the update logic.

class CustomModel(keras.Model):
    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss   = self.compute_loss(y=y, y_pred=y_pred)

        grads = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))

        # Keras 3: update metrics from compile() by hand
        for metric in self.metrics:
            if metric.name == 'loss':
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred)

        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        x, y = data
        y_pred = self(x, training=False)
        loss   = self.compute_loss(y=y, y_pred=y_pred)
        for metric in self.metrics:
            if metric.name == 'loss':
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred)
        return {m.name: m.result() for m in self.metrics}

Legacy (Keras 2 / TF ≤ 2.15): self.compiled_loss(y, y_pred) + self.compiled_metrics.update_state(y, y_pred). Both were removed in Keras 3 — iterate self.metrics as above.

Learning Rate Schedules

# Constant (default)
optimizer = keras.optimizers.Adam(learning_rate=1e-3)

# Step decay
schedule = keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate=1e-3,
    decay_steps=1000,
    decay_rate=0.96,
    staircase=True,     # floor step count to integer multiples of decay_steps
)
optimizer = keras.optimizers.Adam(learning_rate=schedule)

# Cosine decay
schedule = keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=1e-2,
    decay_steps=5000,
    alpha=0.0,          # minimum lr = alpha * initial_lr
)

# Polynomial decay
keras.optimizers.schedules.PolynomialDecay(1e-2, 5000, end_learning_rate=1e-4, power=0.5)

# Piecewise constant
keras.optimizers.schedules.PiecewiseConstantDecay(
    boundaries=[3000, 6000],
    values=[1e-2, 1e-3, 1e-4],
)

# Inverse time
keras.optimizers.schedules.InverseTimeDecay(1e-2, decay_steps=1000, decay_rate=1.0)

Gradient Clipping

# Clip by global norm (recommended — preserves direction)
optimizer = keras.optimizers.Adam(learning_rate=1e-3, clipnorm=1.0)

# Clip by value
optimizer = keras.optimizers.Adam(learning_rate=1e-3, clipvalue=0.5)

# Manual (in custom loop)
grads, _ = tf.clip_by_global_norm(grads, clip_norm=1.0)

Class Weights and Sample Weights

# Imbalanced classes
class_weight = {0: 1.0, 1: 10.0}   # up-weight minority class
model.fit(x, y, class_weight=class_weight)

# Per-sample weights (e.g., importance sampling)
import numpy as np
sample_weight = np.ones(len(y))
sample_weight[y == 1] = 5.0
model.fit(x, y, sample_weight=sample_weight)

Distributed Training

# MirroredStrategy — single machine, multiple GPUs
strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    model = build_model()
    model.compile(optimizer='adam', loss='mse')

model.fit(dataset, epochs=10)

# MultiWorkerMirroredStrategy — multi-machine
strategy = tf.distribute.MultiWorkerMirroredStrategy()

# TPU
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)

Always create the model and compile inside strategy.scope(). Dataset batching should yield global_batch_size = per_replica_batch * num_replicas.