Cheatsheet

import tensorflow as tf
import numpy as np

# ===== Tensors =====
x = tf.constant([1, 2, 3], dtype=tf.float32)
y = tf.Variable(0.0)
z = tf.zeros([3, 4]); z2 = tf.random.normal([3, 4])
x.shape; x.dtype; x.numpy()                # convert back to NumPy

# Type casting, reshape, broadcasting all behave like NumPy
tf.reshape(z, [4, 3])
tf.cast(x, tf.float16)

# Device placement
with tf.device('/GPU:0'):
    a = tf.matmul(z, z2, transpose_b=True)

# ===== tf.data pipelines =====
ds = (tf.data.Dataset
        .from_tensor_slices((X_train, y_train))
        .shuffle(10_000)
        .batch(64)
        .prefetch(tf.data.AUTOTUNE))

# Disk-backed pipelines
files = tf.data.Dataset.list_files('/data/train/*.tfrecord')
ds = files.interleave(tf.data.TFRecordDataset, cycle_length=4).map(parse_fn).batch(64).prefetch(tf.data.AUTOTUNE)

# Image-folder shortcut
img_ds = tf.keras.utils.image_dataset_from_directory(
    '/data/train', image_size=(224, 224), batch_size=32, label_mode='int')

# ===== Models =====
# Sequential
seq = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(32,)),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10),
])

# Functional (multi-input/output)
inputs = tf.keras.Input(shape=(32,))
x_ = tf.keras.layers.Dense(64, activation='relu')(inputs)
out = tf.keras.layers.Dense(10)(x_)
fn_model = tf.keras.Model(inputs, out)

# Subclassed (custom step)
class Net(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.d1 = tf.keras.layers.Dense(64, activation='relu')
        self.d2 = tf.keras.layers.Dense(10)
    def call(self, x):
        return self.d2(self.d1(x))

# ===== Compile + fit =====
seq.compile(
    optimizer=tf.keras.optimizers.AdamW(1e-3),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy'],
)
seq.fit(ds, validation_data=val_ds, epochs=10, callbacks=[
    tf.keras.callbacks.EarlyStopping(patience=3, restore_best_weights=True),
    tf.keras.callbacks.ModelCheckpoint('best.keras', save_best_only=True),
    tf.keras.callbacks.TensorBoard('logs/run'),
    tf.keras.callbacks.ReduceLROnPlateau(patience=2),
])

# ===== Evaluate + predict =====
seq.evaluate(test_ds)
preds = seq.predict(test_ds)
classes = preds.argmax(axis=-1)

# ===== Save / load =====
seq.save('model.keras')                    # full model
loaded = tf.keras.models.load_model('model.keras')
seq.save_weights('weights/ckpt')           # just weights
seq.load_weights('weights/ckpt')

# SavedModel (for TF Serving / TF Hub)
seq.export('saved_model_dir')              # exports the inference signature

# ===== Custom training loop (when fit is not enough) =====
opt = tf.keras.optimizers.Adam(1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
@tf.function
def step(xb, yb, m):
    with tf.GradientTape() as tape:
        logits = m(xb, training=True); loss = loss_fn(yb, logits)
    grads = tape.gradient(loss, m.trainable_weights)
    opt.apply_gradients(zip(grads, m.trainable_weights))
    return loss

# ===== Conversion =====
# TF Lite (mobile)
conv = tf.lite.TFLiteConverter.from_keras_model(seq)
conv.optimizations = [tf.lite.Optimize.DEFAULT]
open('model.tflite', 'wb').write(conv.convert())

# TF.js (browser)
# tensorflowjs_converter --input_format=keras model.keras tfjs_model/

# ONNX (interop)
# pip install tf2onnx
# python -m tf2onnx.convert --saved-model saved_model_dir --output model.onnx

# ===== Performance =====
# - tf.data.AUTOTUNE prefetch/interleave/map: keep GPU fed
# - mixed precision: tf.keras.mixed_precision.set_global_policy('mixed_float16')
# - jit_compile=True on tf.function for fused ops (where supported)
# - tf.distribute.MirroredStrategy for multi-GPU; nothing else changes

# ===== Pitfalls =====
# - SparseCategoricalCrossentropy(from_logits=True) — match output: no softmax in last layer
# - Forgetting tf.AUTOTUNE on the pipeline: GPU sits idle
# - Saving .h5 in 2026 (Keras 3 default is .keras format)
# - Mixing tf.keras (TF) with standalone keras package — pick one
# - Calling .numpy() inside a tight training loop (synchronous → kills throughput)

# tf.data | Sequential | compile fit evaluate | callbacks | save export