TensorFlow Cheatsheet

Inference

Use this TensorFlow reference while you build software engineering projects, review code, or refresh the syntax you reach for most.

model.predict() vs model(x)

MethodReturnsGradientsBatchingtraining flag
model(x)tf.TensorYes (inside GradientTape)ManualPass explicitly
model.predict(x)NumPy arrayNoAutomaticAlways False
model.predict_on_batch(x)NumPy arrayNoNone (single batch)Always False
# model(x) — eager call, returns tensor
logits = model(x_tensor, training=False)
probs  = tf.nn.softmax(logits)

# model.predict() — handles large datasets, returns numpy
preds = model.predict(x_test, batch_size=64, verbose=1)

# model.predict() with a tf.data.Dataset
preds = model.predict(test_dataset, steps=len(test_dataset))

Batch Inference

import tensorflow as tf
import numpy as np

# numpy input — predict() splits into batches automatically
preds = model.predict(x_test, batch_size=128)

# tf.data pipeline — recommended for large datasets
test_ds = (
    tf.data.Dataset.from_tensor_slices(x_test)
    .batch(128)
    .prefetch(tf.data.AUTOTUNE)
)
preds = model.predict(test_ds)

# Manual loop (custom postprocessing per batch)
results = []
for batch in test_ds:
    out = model(batch, training=False)
    results.append(out.numpy())
preds = np.concatenate(results, axis=0)

Postprocessing Predictions

# Classification — get class index
class_ids = np.argmax(preds, axis=1)

# Classification — confidence score
confidence = np.max(preds, axis=1)

# Top-k
top_k_vals, top_k_idx = tf.math.top_k(preds, k=5)

# Binary threshold
binary_preds = (preds[:, 0] > 0.5).astype(int)

# Threshold sweep (find optimal threshold)
from sklearn.metrics import f1_score
thresholds = np.linspace(0, 1, 100)
f1s = [f1_score(y_true, preds > t) for t in thresholds]
best_t = thresholds[np.argmax(f1s)]

Intermediate Layer Outputs

# Functional / Sequential: build a feature extractor
feature_extractor = tf.keras.Model(
    inputs=model.inputs,
    outputs=model.get_layer('conv2d_2').output,
)
features = feature_extractor(x, training=False)

# Multiple intermediate outputs
layer_names = ['block1_conv2', 'block2_conv2', 'block3_conv3']
outputs = [model.get_layer(name).output for name in layer_names]
multi_output_model = tf.keras.Model(inputs=model.inputs, outputs=outputs)
activations = multi_output_model(x, training=False)

# GradCAM-style: need both prediction and conv output
grad_model = tf.keras.Model(
    inputs=model.inputs,
    outputs=[model.get_layer('last_conv').output, model.output],
)
with tf.GradientTape() as tape:
    conv_out, preds = grad_model(x)
    target_class_score = preds[:, class_id]
grads = tape.gradient(target_class_score, conv_out)

tf.function for Fast Inference

@tf.function(input_signature=[tf.TensorSpec([None, 224, 224, 3], tf.float32)])
def predict_fn(x):
    return model(x, training=False)

# First call traces; subsequent calls use compiled graph
output = predict_fn(x_batch)

TensorFlow Lite Inference

import tensorflow as tf

interpreter = tf.lite.Interpreter(model_path='model.tflite')
interpreter.allocate_tensors()

in_details  = interpreter.get_input_details()
out_details = interpreter.get_output_details()

# Set input (must match dtype and shape exactly)
interpreter.set_tensor(in_details[0]['index'], x_sample.astype('float32'))
interpreter.invoke()
output = interpreter.get_tensor(out_details[0]['index'])

# Resize input tensor for variable batch size
interpreter.resize_input_tensor(in_details[0]['index'], [batch, 224, 224, 3])
interpreter.allocate_tensors()

TF Serving REST API

import requests, json, numpy as np

payload = json.dumps({'instances': x_test[:10].tolist()})
response = requests.post(
    'http://localhost:8501/v1/models/my_model:predict',
    data=payload,
    headers={'Content-Type': 'application/json'},
)
preds = np.array(response.json()['predictions'])

Calibration and Uncertainty

# Temperature scaling (post-hoc calibration)
temperature = 1.5

def calibrated_predict(x):
    logits = model(x, training=False)      # raw logits (no softmax output)
    return tf.nn.softmax(logits / temperature)

# Monte Carlo Dropout (Bayesian approximation)
n_passes = 50
preds = tf.stack([model(x, training=True) for _ in range(n_passes)], axis=0)
mean_pred = tf.reduce_mean(preds, axis=0)
std_pred  = tf.math.reduce_std(preds, axis=0)

Streaming / Online Inference

# Stateful RNN inference (one step at a time)
model = tf.keras.Sequential([
    tf.keras.Input(batch_shape=(1, 1, features)),
    tf.keras.layers.LSTM(64, stateful=True),
    tf.keras.layers.Dense(1),
])

model.layers[0].reset_state()  # reset LSTM state between sequences
for timestep in sequence:
    x = timestep[np.newaxis, np.newaxis, :]  # (1, 1, features)
    pred = model.predict(x)

Performance Tips

TipImpact
@tf.function on predict loopAvoids per-call Python overhead
model.predict() > manual loop for large dataHandles batching, reduces Python GIL contention
training=False in model(x)Disables dropout, uses BN inference stats
Batch size as large as VRAM allowsMaximizes GPU utilization
tf.data.AUTOTUNE prefetchOverlaps data loading with inference
TFLite + INT8 quantization4x smaller, 2-4x faster on CPU/edge
tf.saved_model.load + signature callAvoids Keras overhead in production
# Production serving pattern: load once, reuse
loaded   = tf.saved_model.load('saved_model_dir/')
infer_fn = loaded.signatures['serving_default']

def serve(x_numpy):
    x_tensor = tf.constant(x_numpy, dtype=tf.float32)
    result   = infer_fn(x_tensor)
    return result['output_0'].numpy()

Exporting a Serving Signature

@tf.function(input_signature=[tf.TensorSpec([None, 224, 224, 3], tf.float32)])
def serve(x):
    return {'predictions': model(x, training=False)}

tf.saved_model.save(
    model,
    'export/',
    signatures={'serving_default': serve},
)