TensorFlow Cheatsheet

Inference

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.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},
)