TensorFlow cheat sheet

A scannable TensorFlow reference: 11 short snippets across 6 topics, each linking back to the lesson it came from.

At a glance

Quick snippets

Tensors and eager execution

Reshaping and indexing

t = tf.reshape(tf.range(24), (2, 3, 4))    # (2, 3, 4)
t[0]                 # (3, 4)
t[:, 1, :]           # (2, 4)
t[..., 0]            # (2, 3)  — ellipsis covers any leading axes

tf.transpose(t, perm=[0, 2, 1]).shape      # (2, 4, 3)
tf.expand_dims(t, axis=-1).shape           # (2, 3, 4, 1)
tf.squeeze(tf.zeros((1, 4, 1))).shape      # (4,)
tf.concat([tf.zeros((2, 3)), tf.ones((2, 3))], axis=0).shape   # (4, 3)
tf.stack([tf.zeros((4,)), tf.ones((4,))], axis=1).shape        # (4, 2)

Reshaping and indexing

# a variable used in a tf.function with a changing Python argument retraces
@tf.function
def scale(x, factor):
    return x * factor

scale(tf.ones((2, 2)), 2.0)          # trace 1 (int vs float matters)
scale(tf.ones((2, 2)), 3)            # trace 2: different Python type
scale(tf.ones((2, 2)), tf.constant(2.0))   # trace 3: tensor, one trace for all values

Full lesson: Tensors and eager execution →

Building a Keras model

Two ways to build

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# Sequential: a linear stack, fine for most feed-forward models
model = keras.Sequential([
    keras.Input(shape=(784,)),                 # declare the shape once, here
    layers.Dense(128, activation="relu"),
    layers.Dropout(0.2),
    layers.Dense(10),                          # logits, no activation
])
model.summary()

Two ways to build

# Functional: branches, multiple inputs or outputs, explicit shapes
inputs = keras.Input(shape=(784,), name="pixels")
x = layers.Dense(128, activation="relu")(inputs)
x = layers.Dropout(0.2)(x)
x = layers.Dense(64, activation="relu")(x)
outputs = layers.Dense(10, name="logits")(x)

model = keras.Model(inputs=inputs, outputs=outputs, name="mlp")
keras.utils.plot_model(model, show_shapes=True)   # needs graphviz

Output layer, loss and compile

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=1e-3),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[keras.metrics.SparseCategoricalAccuracy(name="acc")],
)

# binary classification: one output, logit, from_logits=True
# loss=keras.losses.BinaryCrossentropy(from_logits=True)

# regression: one output, linear activation, MSE or Huber
# loss=keras.losses.MeanSquaredError()

Full lesson: Building a Keras model →

Training, saving and serving

Saving and serving

model.save("model.keras")                # Keras v3 format: architecture, weights, state
model.save("saved_model_dir")            # SavedModel, for TF Serving

restored = keras.models.load_model("model.keras")
restored.predict(x_test[:8])

latest = keras.models.load_model("best.keras")   # written by ModelCheckpoint
latest.evaluate(x_test, y_test)                  # evaluate, not predict, for metrics

keras.backend.clear_session()            # free GPU memory between experiments

Saving and serving

# put preprocessing INSIDE the model so serving matches training exactly
normalizer = keras.layers.Normalization(axis=-1)
normalizer.adapt(x_train)                # learns mean and variance from train only

model = keras.Sequential([
    keras.Input(shape=x_train.shape[1:]),
    normalizer,
    keras.layers.Dense(64, activation="relu"),
    keras.layers.Dense(1, activation="sigmoid"),
])
model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])

Full lesson: Training, saving and serving →

Text and image preprocessing layers

Why preprocessing belongs in the model

import tensorflow as tf

vectoriser = tf.keras.layers.TextVectorization(
    max_tokens=20_000,        # vocabulary size, OOV index 0 is prepended
    output_mode="int",
    output_sequence_length=64,
    standardize="lower_and_strip_punctuation",
    ngrams=None,
)
vectoriser.adapt(train_texts)      # learns the vocabulary from training data only
print(vectoriser.get_vocabulary()[:10])
print(vectoriser(["The quick brown fox"]).shape)     # (1, 64)

Full lesson: Text and image preprocessing layers →

TensorBoard and experiment tracking

Logging during training

tensorboard --logdir logs/fit --port 6006
# compare several runs at once: point logdir at the parent directory
tensorboard --logdir logs/

Custom scalars, images and text

# a custom callback writing metrics TensorBoard does not track automatically
class LRLogger(tf.keras.callbacks.Callback):
    def __init__(self, log_dir):
        super().__init__()
        self.writer = tf.summary.create_file_writer(log_dir)

    def on_epoch_end(self, epoch, logs=None):
        lr = float(tf.keras.backend.get_value(self.model.optimizer.learning_rate))
        with self.writer.as_default():
            tf.summary.scalar("train/learning_rate", lr, step=epoch)
        self.writer.flush()

Full lesson: TensorBoard and experiment tracking →

Debugging, profiling and export formats

Profiling a training step

# the profiler records a short window; open the result in TensorBoard
tf.profiler.experimental.start("logs/profile")

for step, (x, y) in enumerate(train_ds.take(20)):
    with tf.profiler.experimental.Trace("train", step_num=step):
        loss = train_step(x, y)
    if step == 10:
        break

tf.profiler.experimental.stop()   # then: tensorboard --logdir logs/profile

Full lesson: Debugging, profiling and export formats →

FAQ

Related cheat sheets

AI Basics AI Agents Math for AI Machine Learning scikit-learn PyTorch

Last refreshed 2026-09-27.