Getting Started
Installation & Basics
TensorFlow 2.x runs eagerly by default (operations execute immediately, like normal Python). Use tf.config.list_physical_devices('GPU') to verify CUDA GPUs are visible. For GPU support install the [and-cuda] extra on Linux/Windows, or use tensorflow-metal on Apple Silicon. TF 2.x merged Keras as tf.keras — the high-level API used in most workflows.
# install TensorFlow (CPU)
pip install tensorflow
# with GPU support on Linux/Windows (CUDA-enabled)
pip install tensorflow[and-cuda]
import tensorflow as tf
# check version and devices
print(tf.__version__)
print("GPU:", tf.config.list_physical_devices('GPU'))
# eager execution is on by default in TF2
print(tf.executing_eagerly()) # TrueFirst Tensors
tf.constant creates an immutable tensor; tf.zeros / tf.ones / tf.random.* create common patterns. Every tensor has a shape, dtype, and rank. .numpy() bridges back to NumPy — on GPU this may trigger a sync copy. Prefer explicit dtypes (float32 is the default for ML) to avoid silent upcasts that break performance.
import tensorflow as tf
# constants are immutable
x = tf.constant([[1, 2], [3, 4]])
y = tf.zeros((3, 3))
z = tf.random.normal((2, 3), mean=0.0, stddev=1.0)
print(x.shape, x.dtype) # (2, 2) int32
print(tf.reduce_sum(x)) # tf.Tensor(10, shape=(), dtype=int32)
# numpy interop is zero-copy when possible
print(x.numpy()) # array([[1, 2], [3, 4]])Eager Execution & Gradients
Eager execution makes TF debuggable like normal Python — print, pdb, and tracebacks all work. GradientTape records forward ops on watched tensors (tf.Variable is watched automatically) so tape.gradient can compute derivatives. This is the foundation of custom training loops. tapes are single-use by default; use persistent=True to call gradient multiple times.
import tensorflow as tf
a = tf.constant(3.0)
b = tf.constant(4.0)
print(a * b + 1) # tf.Tensor(13.0, ...)
# GradientTape records ops for automatic differentiation
with tf.GradientTape() as tape:
tape.watch(a)
y = a ** 2 + b # y = a^2 + b, dy/da = 2a
grad = tape.gradient(y, a)
print(grad) # tf.Tensor(6.0, ...) (= 2*3)tf.function & Graphs
@tf.function converts a Python function into a portable, optimizable TensorFlow graph — much faster on GPU and required for distribution. The function is traced once per input signature, so avoid Python side effects or data-dependent Python control flow inside (use tf.cond / tf.while_loop or the function will retrace). Variables must be created once, outside the function, or reuse tf.Variable with the same name.
import tensorflow as tf
@tf.function # traces a graph; runs in graph mode after first call
def train_step(x, y):
with tf.GradientTape() as tape:
preds = x * w + b
loss = tf.reduce_mean((preds - y) ** 2)
grads = tape.gradient(loss, [w, b])
optimizer.apply_gradients(zip(grads, [w, b]))
return loss
w = tf.Variable(0.1)
b = tf.Variable(0.0)
optimizer = tf.keras.optimizers.SGD(0.01)
# first call builds the graph (slower); subsequent calls are fast
loss = train_step(tf.constant([1., 2.]), tf.constant([2., 4.]))Mixed Precision
Mixed precision computes in float16 but keeps float32 master copies, giving 2-3x throughput and halving memory on Tensor Core GPUs (Volta+). The 'mixed_float16' policy handles loss scaling automatically to prevent underflow. Force the output layer to float32 for numerical stability. Needs compute capability >= 7.0; on older GPUs it may be slower.
import tensorflow as tf
# enable mixed precision: compute in float16, store in float32
tf.keras.mixed_precision.set_global_policy('mixed_float16')
model = tf.keras.Sequential([
tf.keras.layers.Dense(64, activation='relu', dtype='float32'),
tf.keras.layers.Dense(10, activation='softmax', dtype='float32')
])
# loss scaling is automatic with mixed_float16 policy
print(tf.keras.mixed_precision.global_policy())Tensors
tf.constant
tf.constant creates an immutable tensor — you cannot assign to elements. For trainable tensors use tf.Variable. tf.random.normal/uniform accept a seed argument for reproducibility; tf.random.set_seed(42) makes the whole program reproducible. Use shape tuples (not ints) for multi-dim creation to avoid surprises.
import tensorflow as tf
a = tf.constant([1, 2, 3])
b = tf.constant([[1.0, 2.0], [3.0, 4.0]], dtype=tf.float32)
zeros = tf.zeros((2, 3))
ones = tf.ones((3, 3))
eye = tf.eye(3)
full = tf.fill((2, 3), 7.0)
arange = tf.range(0, 10, 2) # [0, 2, 4, 6, 8]
linspace = tf.linspace(0.0, 1.0, 5)
rand = tf.random.uniform((2, 3))
randn = tf.random.normal((2, 3))
# constants are immutable: a[0] = 5 -> TypeErrortf.Variable
tf.Variable is a mutable tensor wrapper used for model parameters and optimizer state. assign / assign_add / assign_sub update it in place — ordinary Python assignment (w = w + 1) creates a new tensor and does not mutate. Variables are watched by GradientTape automatically, unlike tf.constant which you must tape.watch() explicitly. Keep Variables on the same device as the ops that use them.
import tensorflow as tf
w = tf.Variable(0.1, dtype=tf.float32)
W = tf.Variable(tf.random.normal((3, 4)))
b = tf.Variable(tf.zeros((4,)))
# read value
print(w.numpy()) # 0.1
# update value
w.assign(0.5)
w.assign_add(0.1) # w += 0.1
w.assign_sub(0.05)
# assign in place on a slice
W[0, 0].assign(99.0)
# watched by GradientTape automatically
with tf.GradientTape() as tape:
y = w ** 2
grad = tape.gradient(y, w) # 2wTensor Operations
Most math mirrors NumPy: tf.reduce_sum / reduce_mean / reduce_max correspond to np.sum/mean/max. Use axis= to reduce a specific dimension and keepdims=True to preserve rank for broadcasting. tf.matmul and the @ operator do matrix multiply (batched when inputs are 3D+). reshape returns a view when possible; transpose always returns a copy in TF.
import tensorflow as tf
a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
b = tf.constant([[5.0, 6.0], [7.0, 8.0]])
a + b; a - b; a * b; a / b # element-wise
c = tf.matmul(a, b) # matrix multiply
c = a @ b # equivalent
tf.reduce_sum(a) # sum of all
tf.reduce_sum(a, axis=0) # sum along axis 0
tf.reduce_mean(a, axis=1) # mean along axis 1
tf.reduce_max(a, axis=1)
tf.reduce_sum(a, axis=1, keepdims=True)
x = tf.range(12)
y = tf.reshape(x, (3, 4))
t = tf.transpose(y) # shape (4, 3)Indexing & Slicing
TF supports NumPy-style slicing in eager mode and inside tf.function. tf.boolean_mask extracts elements where the mask is True (flattened). tf.gather selects along axis 0 by default (like fancy indexing). tf.where(cond, x, y) is the vectorized ternary — when cond is True pick from x else y. For assignment use tf.tensor_scatter_nd_update or a tf.Variable.
import tensorflow as tf
x = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
x[0] # first row
x[1, 2] # element (1,2) -> 6
x[:, 1] # second column
x[0:2] # first two rows
x[:, ::2] # every other column
# boolean mask
mask = x > 4
sel = tf.boolean_mask(x, mask) # tensor([5, 6, 7, 8, 9])
# gather by indices
idx = tf.constant([0, 2])
rows = tf.gather(x, idx) # rows 0 and 2
# where: select from two tensors
tf.where(x > 4, x, tf.zeros_like(x))Ragged & Sparse Tensors
Ragged tensors store variable-length sequences without padding — perfect for NLP batches of different lengths. Sparse tensors store only nonzero indices and values, saving memory for highly sparse data (recommendations, graphs). Most ops have sparse-aware variants; convert to dense only when an op requires it. tf.RaggedTensor.to_tensor() pads to a rectangle for embedding lookups.
import tensorflow as tf
# ragged: variable-length sequences
ragged = tf.ragged.constant([[1, 2, 3], [4], [5, 6]])
print(ragged.shape) # (3, None)
ragged.to_tensor() # pad to (3, 3)
# sparse: mostly-zero tensors
indices = [[0, 0], [1, 2], [2, 1]]
values = [1.0, 2.0, 3.0]
sparse = tf.sparse.SparseTensor(
indices=indices, values=values, dense_shape=[3, 3])
dense = tf.sparse.to_dense(sparse)
sp_sum = tf.sparse.reduce_sum(sparse)
sp_matmul = tf.sparse.sparse_dense_matmul(sparse, tf.eye(3))Broadcasting & Math
Broadcasting stretches smaller-dim tensors to match larger ones, element-wise. clip_by_value is the TF clamp. tf.maximum/minimum are element-wise (different from tf.reduce_max/min which reduce a dimension). cumsum/cumprod produce running sums/products along an axis — useful for cumulative metrics and prefix-sum algorithms in custom layers.
import tensorflow as tf
# broadcasting follows NumPy rules
a = tf.constant([[1.0, 2.0, 3.0]]) # shape (1, 3)
b = tf.constant([[10.0], [20.0]]) # shape (2, 1)
c = a + b # shape (2, 3)
tf.abs(x); tf.sqrt(x); tf.exp(x); tf.log(x)
tf.square(x); tf.pow(x, 3)
tf.maximum(a, b); tf.minimum(a, b)
tf.clip_by_value(x, 0.0, 1.0)
tf.norm(x) # L2 norm
tf.norm(x, ord=1) # L1 norm
tf.cumsum(tf.range(5)) # [0, 1, 3, 6, 10]
tf.cumprod(tf.range(1, 5)) # [1, 2, 6, 24]Keras Models
Sequential API
Sequential is the simplest Keras model — a linear pipeline where each layer's output feeds the next. input_shape excludes the batch dimension (use None for variable batch). model.summary() prints layer shapes and parameter counts. If you don't pass input_shape, the model builds lazily on the first fit/forward call — call build() explicitly to inspect it earlier.
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax'),
])
model.summary()
# build later with input_shape (or first call)
model = tf.keras.Sequential()
model.add(tf.keras.layers.Dense(128, activation='relu'))
model.add(tf.keras.layers.Dense(10))
model.build((None, 784)) # None = variable batchFunctional API
The Functional API lets layers branch, merge, and skip — required for anything non-linear (ResNets, multi-input/output, siamese nets). Define Input tensors, chain layers by calling them, then pass inputs/outputs to Model. The resulting model behaves identically to a Sequential one for fit/predict. Name inputs/outputs for clean saving and serving signatures.
import tensorflow as tf
inputs = tf.keras.Input(shape=(784,), name='img')
x = tf.keras.layers.Dense(128, activation='relu')(inputs)
x = tf.keras.layers.Dropout(0.5)(x)
x = tf.keras.layers.Dense(64, activation='relu')(x)
outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs, name='mlp')
# multi-input / multi-output
img_in = tf.keras.Input(shape=(224, 224, 3))
meta_in = tf.keras.Input(shape=(10,))
x = tf.keras.layers.Conv2D(32, 3, activation='relu')(img_in)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.concatenate([x, meta_in])
out = tf.keras.layers.Dense(1, activation='sigmoid')(x)
model = tf.keras.Model([img_in, meta_in], out)Model Subclassing
Subclassing tf.keras.Model gives full control of forward (like PyTorch). Define layers in __init__, implement call(inputs, training=). The training flag toggles dropout and BatchNorm behavior. Subclassing is more flexible but loses some Functional-API benefits: harder to inspect/plot, no automatic input shape inference, and serialization is trickier. Prefer Functional for most architectures; subclass for research models.
import tensorflow as tf
class MLP(tf.keras.Model):
def __init__(self, hidden=128, num_classes=10):
super().__init__()
self.dense1 = tf.keras.layers.Dense(hidden, activation='relu')
self.drop = tf.keras.layers.Dropout(0.5)
self.dense2 = tf.keras.layers.Dense(hidden, activation='relu')
self.out = tf.keras.layers.Dense(num_classes, activation='softmax')
def call(self, inputs, training=False):
x = self.dense1(inputs)
x = self.drop(x, training=training)
x = self.dense2(x)
return self.out(x)
model = MLP(hidden=64)
model.build((None, 784))
model.summary()Model Compilation
compile wires the optimizer, loss, and metrics together before fit. Strings ('adam', 'mse') are shortcuts to the corresponding classes; pass objects for custom hyperparameters. For multi-output models pass dictionaries mapping output names to losses. run_eagerly=True disables tf.function tracing so you can use pdb and print — invaluable for debugging, but slow.
model.compile(
optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'],
loss_weights={'cls': 1.0, 'box': 0.5},
)
# explicit objects (more control)
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
loss=tf.keras.losses.SparseCategoricalCrossentropy(),
metrics=[tf.keras.metrics.SparseCategoricalAccuracy(name='acc')],
run_eagerly=False, # debug mode: run without graph
)Model Summary & Plot
summary() prints a per-layer table with output shapes and parameter counts — the fastest way to spot shape bugs. plot_model saves a visual graph (requires pydot and Graphviz). get_config returns a serializable dict so the architecture can be rebuilt without weights — used by model.to_json() and the Keras save format. For Functional models, config fully captures the architecture; for subclassed models you must implement get_config yourself.
model.summary()
print(model.layers)
print(model.inputs, model.outputs)
# plot architecture (needs pydot + graphviz)
tf.keras.utils.plot_model(
model, to_file='model.png', show_shapes=True,
show_layer_names=True, rankdir='TB',
)
print(f"{model.count_params():,} params")
config = model.get_config()
new_model = tf.keras.Model.from_config(config)Preprocessing in Model
Putting preprocessing inside the model (Rescaling, augmentation Keras layers) makes it deployable: the same SavedModel accepts raw images at inference, no separate preprocessing step. Keras preprocessing layers run on GPU and inside tf.function, so they're fast. For training-only augmentation, set training=True in fit so the layers are inactive at inference. Rescaling(1./255) converts uint8 [0,255] images to float [0,1].
import tensorflow as tf
# preprocessing as part of the model (deployable)
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(224, 224, 3)),
tf.keras.layers.Rescaling(1./255),
tf.keras.layers.RandomFlip('horizontal'),
tf.keras.layers.RandomRotation(0.1),
tf.keras.layers.Conv2D(32, 3, activation='relu'),
# ...
])
# the model now includes preprocessing, so inference needs raw imagesLayers
Dense Layer
Dense is the fully-connected layer: y = activation(xW + b). Weight shape is (in_features, out_features) — the transpose of PyTorch. kernel_initializer defaults to Glorot/Xavier uniform, good for tanh/sigmoid; use 'he_normal' for ReLU-family. Add L2 regularization via kernel_regularizer. GELU is the default activation in transformers.
import tensorflow as tf
dense = tf.keras.layers.Dense(
units=64, activation='relu', use_bias=True,
kernel_initializer='glorot_uniform',
bias_initializer='zeros',
kernel_regularizer=tf.keras.regularizers.l2(1e-4),
)
x = tf.random.normal((4, 128))
y = dense(x) # shape (4, 64)
print(dense.kernel.shape) # (128, 64)
print(dense.bias.shape) # (64,)
dense = tf.keras.layers.Dense(64, activation=tf.nn.gelu)Conv2D Layer
TF/Keras uses NHWC (channels-last) by default, opposite to PyTorch's NCHW. padding='same' pads so output spatial size = input size / stride; 'valid' means no padding. DepthwiseConv2D applies one filter per input channel (used in MobileNet). Conv2DTranspose upsamples (used in segmentation decoders). The kernel shape is (kH, kW, in_channels, filters).
import tensorflow as tf
conv = tf.keras.layers.Conv2D(
filters=32, kernel_size=3, strides=(1, 1),
padding='same', activation='relu', use_bias=True,
)
x = tf.random.normal((4, 28, 28, 3)) # NHWC
out = conv(x)
print(out.shape) # (4, 28, 28, 32) with padding='same'
dw = tf.keras.layers.DepthwiseConv2D(3, padding='same', activation='relu')
up = tf.keras.layers.Conv2DTranspose(32, 3, strides=2, padding='same')Recurrent Layers (LSTM/GRU)
Input shape for RNNs is (batch, timesteps, features). return_sequences=True is needed when stacking RNNs or for sequence-to-sequence models; False for sequence-to-vector. Bidirectional doubles the output size by concatenating forward and backward. recurrent_dropout is convenient but slow (no cuDNN acceleration) — prefer plain dropout on inputs instead.
import tensorflow as tf
lstm = tf.keras.layers.LSTM(units=128, return_sequences=False)
x = tf.random.normal((4, 10, 64)) # (batch, time, features)
out = lstm(x) # (4, 128)
# stacked LSTM needs return_sequences=True
model = tf.keras.Sequential([
tf.keras.layers.LSTM(128, return_sequences=True, input_shape=(10, 64)),
tf.keras.layers.LSTM(64),
tf.keras.layers.Dense(10),
])
gru = tf.keras.layers.GRU(128)
bi = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64))
out = bi(x) # shape (4, 128) — 64 * 2Dropout & Regularization
Dropout randomly zeros units during training to prevent co-adaptation — it's automatically disabled at inference (training=False). SpatialDropout2D drops whole feature maps, better for conv nets. L2 (weight decay) is the most common regularizer. LayerNormalization normalizes over the feature axis (per-sample), unlike BatchNorm which normalizes over the batch axis — preferred for transformers and RNNs.
import tensorflow as tf
drop = tf.keras.layers.Dropout(0.5)
x = tf.random.normal((4, 10))
y = drop(x, training=True) # training=True activates dropout
tf.keras.layers.SpatialDropout2D(0.5) # drop entire feature maps
tf.keras.layers.GaussianDropout(0.5)
tf.keras.layers.AlphaDropout(0.5)
# weight regularization
tf.keras.layers.Dense(64,
kernel_regularizer=tf.keras.regularizers.l2(1e-4),
activity_regularizer=tf.keras.regularizers.l1(1e-5))
norm = tf.keras.layers.LayerNormalization()Batch Normalization
BatchNormalization normalizes activations per channel, keeping EMA running stats for inference. axis=-1 normalizes the last (channel) axis for NHWC. Put BN after the conv and before the activation, and disable the conv's bias (use_bias=False) since BN's beta handles the shift. momentum controls how fast running stats track the data — high values (0.99) are more stable. BN behaves differently in train vs inference (training flag).
import tensorflow as tf
bn = tf.keras.layers.BatchNormalization(
axis=-1, momentum=0.99, epsilon=1e-3,
)
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, 3, padding='same', use_bias=False),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.ReLU(),
])
print(len(bn.weights)) # 4
print(len(bn.trainable_weights)) # 2Custom Layer
Subclass Layer to build custom ops. Define weights in build() (called once with the input shape) rather than __init__ so the layer adapts to input size. Implement get_config() so the layer can be serialized and reloaded. add_weight registers the tensor with Keras' tracking (so it shows up in weights, is saved, and is moved to device). For stateless transforms, you can skip build() and just override call().
import tensorflow as tf
class Linear(tf.keras.layers.Layer):
def __init__(self, units=32, **kwargs):
super().__init__(**kwargs)
self.units = units
def build(self, input_shape):
self.w = self.add_weight(
shape=(input_shape[-1], self.units),
initializer='glorot_uniform', trainable=True, name='w')
self.b = self.add_weight(
shape=(self.units,), initializer='zeros', trainable=True, name='b')
def call(self, inputs):
return tf.matmul(inputs, self.w) + self.b
def get_config(self):
config = super().get_config()
config.update({'units': self.units})
return config
layer = Linear(64)
out = layer(tf.random.normal((4, 10)))Losses
BinaryCrossentropy
BinaryCrossentropy is for binary tasks or multi-label (each class independent). Pass from_logits=True if your model outputs raw logits — it's numerically safer than applying sigmoid + BCE separately. pos_weight up-weights positive examples to handle imbalance. For multi-label problems with C classes, target shape is (N, C) of 0/1 and predictions are (N, C) probabilities or logits.
import tensorflow as tf
bce = tf.keras.losses.BinaryCrossentropy(from_logits=False)
y_true = tf.constant([[1.0], [0.0], [1.0], [0.0]])
y_pred = tf.constant([[0.9], [0.1], [0.8], [0.2]])
loss = bce(y_true, y_pred)
# from_logits=True is more numerically stable
bce_logits = tf.keras.losses.BinaryCrossentropy(from_logits=True)
logits = tf.constant([[2.0], [-2.0], [1.5], [-1.5]])
loss = bce_logits(y_true, logits)
bce = tf.keras.losses.BinaryCrossentropy(
from_logits=True, pos_weight=tf.constant([5.0]))CategoricalCrossentropy
CategoricalCrossentropy expects one-hot targets; SparseCategoricalCrossentropy expects integer class indices (more memory-efficient for many classes). Always pass from_logits=True when the model doesn't end in softmax — it uses the log-sum-exp trick for numerical stability. For label smoothing use CategoricalCrossentropy(label_smoothing=0.1), a strong regularizer for classification and NLP.
import tensorflow as tf
# one-hot labels
cce = tf.keras.losses.CategoricalCrossentropy(from_logits=False)
y_true = tf.constant([[0, 1, 0], [1, 0, 0]], dtype=tf.float32)
y_pred = tf.constant([[0.1, 0.8, 0.1], [0.7, 0.2, 0.1]])
loss = cce(y_true, y_pred)
# integer labels -> use sparse version
scce = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
y_true_int = tf.constant([1, 0])
logits = tf.constant([[1.0, 3.0, 0.5], [2.5, 0.2, 0.1]])
loss = scce(y_true_int, logits)MSE & Regression Losses
MSE penalizes large errors quadratically (sensitive to outliers); MAE is linear. Huber behaves like MSE near zero and MAE far away — the default for object detection regression. LogCosh is similar but twice differentiable everywhere, useful when you need second-order gradients. All reduction defaults to mean over the batch; override with reduction=tf.keras.losses.Reduction.NONE for per-sample losses.
import tensorflow as tf
mse = tf.keras.losses.MeanSquaredError()
y_true = tf.constant([3.0, -0.5, 2.0, 7.0])
y_pred = tf.constant([2.5, 0.0, 2.0, 8.0])
loss = mse(y_true, y_pred)
mae = tf.keras.losses.MeanAbsoluteError()
loss = mae(y_true, y_pred)
huber = tf.keras.losses.Huber(delta=1.0)
loss = huber(y_true, y_pred)
logcosh = tf.keras.losses.LogCosh()
loss = logcosh(y_true, y_pred)Custom Loss
A custom loss is any callable taking (y_true, y_pred) and returning a scalar tensor. For serialization or hyperparameters, subclass tf.keras.losses.Loss and implement call() + get_config(). Keep every op differentiable. FocalLoss down-weights easy examples to focus on hard ones — popular for object detection. Use Reduction.NONE in call() and let Keras reduce at the end if you want masking to work.
import tensorflow as tf
def weighted_mse(y_true, y_pred):
weights = tf.cast(y_true != 0, tf.float32)
return tf.reduce_sum(weights * tf.square(y_true - y_pred)) / tf.reduce_sum(weights)
class FocalLoss(tf.keras.losses.Loss):
def __init__(self, alpha=0.25, gamma=2.0, **kwargs):
super().__init__(**kwargs)
self.alpha = alpha
self.gamma = gamma
def call(self, y_true, y_pred):
bce = tf.keras.losses.binary_crossentropy(y_true, y_pred)
p = y_pred
pt = p * y_true + (1 - p) * (1 - y_true)
return self.alpha * tf.pow(1 - pt, self.gamma) * bce
def get_config(self):
config = super().get_config()
config.update({'alpha': self.alpha, 'gamma': self.gamma})
return config
loss = FocalLoss(alpha=0.25, gamma=2.0)Loss from Logits
Computing softmax then cross-entropy separately can overflow for large logits (e^100). Passing from_logits=True lets TF fuse them with the log-sum-exp trick, which is numerically stable. The functional tf.nn.*_cross_entropy_with_logits APIs do the same thing. Make sure your model does NOT end in softmax/sigmoid when using from_logits=True, otherwise you double-apply the activation.
import tensorflow as tf
# why logits? numerical stability
logits = tf.constant([[100.0, 0.0, 0.0]])
# BAD: softmax first then CCE (can overflow)
probs = tf.nn.softmax(logits)
loss_bad = tf.keras.losses.CategoricalCrossentropy()(
tf.constant([[1.0, 0.0, 0.0]]), probs)
# GOOD: from_logits=True uses log-sum-exp internally
loss_good = tf.keras.losses.CategoricalCrossentropy(from_logits=True)(
tf.constant([[1.0, 0.0, 0.0]]), logits)
loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=y, logits=z)
loss = tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=z)Multi-Output Losses
For multi-output models, pass a dict mapping output names to loss functions and a loss_weights dict to balance them. The total loss is the weighted sum. This is the standard pattern for multi-task learning — each head gets its own loss and metric. Pass training targets as a dict with the same keys at fit time. Tune loss_weights so no single task dominates the gradient.
import tensorflow as tf
model.compile(
optimizer='adam',
loss={
'priority': tf.keras.losses.BinaryCrossentropy(from_logits=True),
'department': tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
},
loss_weights={'priority': 1.0, 'department': 0.5},
metrics={'priority': ['accuracy'], 'department': ['accuracy']},
)
inputs = tf.keras.Input(shape=(64,))
x = tf.keras.layers.Dense(64, activation='relu')(inputs)
out_a = tf.keras.layers.Dense(1, name='priority')(x)
out_b = tf.keras.layers.Dense(10, name='department')(x)
model = tf.keras.Model(inputs, [out_a, out_b])
model.fit(x_train, {'priority': y_pri, 'department': y_dep}, epochs=5)Optimizers
Adam
Adam is the default optimizer for most tasks — adaptive per-parameter learning rates, robust to hyperparameter choice. beta_1/beta_2 control the EMA decay of the first and second moments. Pass a LearningRateSchedule object instead of a float to decay LR over training. global_clipnorm rescales all gradients if their combined norm exceeds the threshold — important for RNNs and transformers.
import tensorflow as tf
opt = tf.keras.optimizers.Adam(learning_rate=1e-3)
opt = tf.keras.optimizers.Adam(
learning_rate=1e-3, beta_1=0.9, beta_2=0.999,
epsilon=1e-7, amsgrad=False,
)
# schedule as learning_rate
lr_schedule = tf.keras.optimizers.schedules.CosineDecay(
initial_learning_rate=1e-3, decay_steps=10000)
opt = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
# gradient clipping
opt = tf.keras.optimizers.Adam(learning_rate=1e-3, global_clipnorm=1.0)SGD
SGD + momentum often generalizes better than Adam on large-batch CNN training, at the cost of more LR tuning. momentum=0.9 is standard; Nesterov gives a small extra edge. weight_decay adds decoupled L2 regularization (Keras 3) — older versions use kernel_regularizer on layers instead. SGD needs warmup and LR scheduling to compete with Adam out of the box.
import tensorflow as tf
opt = tf.keras.optimizers.SGD(learning_rate=0.01)
opt = tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.9)
opt = tf.keras.optimizers.SGD(
learning_rate=0.01, momentum=0.9, nesterov=True)
# with weight decay (Keras 2.13+ / 3.x)
opt = tf.keras.optimizers.SGD(
learning_rate=0.01, momentum=0.9, weight_decay=1e-4)
# clip gradients by value
opt = tf.keras.optimizers.SGD(
learning_rate=0.01, momentum=0.9, clipvalue=0.5)AdamW & Other Optimizers
AdamW decouples weight decay from the gradient update, which matters at scale (transformers, large CNNs). RMSprop is the classic RNN optimizer. Adagrad adapts per-parameter but can stall early — fine for sparse features. FTRL is designed for large-scale online learning with sparse features (click-through prediction). For new projects, start with AdamW.
import tensorflow as tf
opt = tf.keras.optimizers.AdamW(
learning_rate=3e-4, weight_decay=0.01,
beta_1=0.9, beta_2=0.999)
opt = tf.keras.optimizers.RMSprop(learning_rate=1e-3, rho=0.9)
opt = tf.keras.optimizers.Adagrad(learning_rate=1e-2)
opt = tf.keras.optimizers.Adadelta(learning_rate=1.0)
opt = tf.keras.optimizers.Nadam(learning_rate=1e-3)
opt = tf.keras.optimizers.Ftrl(learning_rate=0.1)Applying Gradients
In a custom training loop you compute gradients with GradientTape, then call opt.apply_gradients(zip(grads, vars)). This is the building block of custom training loops (see custom-training-loop section). clip_by_global_norm rescales gradients so their combined L2 norm stays under the threshold. None gradients mean the variable didn't affect the loss — common for unused parameters and worth logging.
import tensorflow as tf
opt = tf.keras.optimizers.Adam(1e-3)
vars = [w, b]
with tf.GradientTape() as tape:
preds = model(x)
loss = loss_fn(y, preds)
grads = tape.gradient(loss, vars)
grads, _ = tf.clip_by_global_norm(grads, 1.0)
opt.apply_gradients(zip(grads, vars))
for v, g in zip(vars, grads):
print(v.name, g is not None, g.numpy().norm() if g is not None else None)Learning Rate Schedules
LearningRateSchedule objects compute LR as a function of the optimizer's step counter. CosineDecay smoothly anneals to alpha*initial_lr — a strong default. CosineDecayRestarts implements SGDR-style warm restarts. Pass the schedule as learning_rate when constructing the optimizer; Keras calls it automatically each step. For warmup, wrap with a custom schedule or use the WarmUp callback from tensorflow/addons.
import tensorflow as tf
sched = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
boundaries=[500, 1500], values=[1e-3, 1e-4, 1e-5])
sched = tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate=1e-3, decay_steps=1000, decay_rate=0.96,
staircase=True)
sched = tf.keras.optimizers.schedules.CosineDecay(
initial_learning_rate=1e-3, decay_steps=10000, alpha=0.0)
opt = tf.keras.optimizers.Adam(learning_rate=sched)
print(opt.learning_rate(opt.iterations))Gradient Clipping & EMA
global_clipnorm (rescales all gradients together) is the most common choice and is what transformers use. clipvalue hard-clamps each element. For custom clipping, modify grads between tape.gradient and apply_gradients. Exponential Moving Average (EMA) of weights often gives a small accuracy bump at inference — apply it after each step and use the averaged weights for evaluation. TF Addons has a MovingAverage optimizer wrapper.
import tensorflow as tf
opt = tf.keras.optimizers.Adam(1e-3, global_clipnorm=1.0)
opt = tf.keras.optimizers.Adam(1e-3, clipvalue=0.5)
opt = tf.keras.optimizers.Adam(1e-3, clipnorm=1.0)
# manual clipping (more control)
grads = tape.gradient(loss, vars)
grads = [tf.clip_by_norm(g, 1.0) if g is not None else g for g in grads]
grads, _ = tf.clip_by_global_norm(grads, 1.0)
opt.apply_gradients(zip(grads, vars))
# exponential moving average of weights (improves generalization)
ema = tf.train.ExponentialMovingAverage(decay=0.999)Training
compile & fit
fit is the high-level training API: it handles batching, shuffling, validation, callbacks, and progress. Pass either validation_data=(x, y) or validation_split=0.2 (held out from training). verbose=2 prints one line per epoch (good for logs); 1 shows a progress bar. The returned history object stores loss and metrics per epoch — plot it to spot overfitting.
import tensorflow as tf
model.compile(
optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'],
)
history = model.fit(
x_train, y_train, batch_size=32, epochs=10,
validation_data=(x_val, y_val),
validation_split=0.2, shuffle=True, verbose=2,
)
print(history.history['loss'])Validation & Metrics
Metrics are stateful objects that accumulate across batches and compute the final value at epoch end — unlike losses they don't affect training. Pass strings ('accuracy') or Metric objects for custom metrics. Subclass tf.keras.metrics.Metric and implement update_state/result/reset_state. reset_state (renamed from reset_states in TF 2.11+) is called at the start of each epoch.
model.compile(
optimizer='adam', loss='sparse_categorical_crossentropy',
metrics=[
'accuracy',
tf.keras.metrics.SparseTopKCategoricalAccuracy(k=5, name='top5_acc'),
tf.keras.metrics.Precision(name='precision'),
tf.keras.metrics.Recall(name='recall'),
],
)
class F1Score(tf.keras.metrics.Metric):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.precision = tf.keras.metrics.Precision()
self.recall = tf.keras.metrics.Recall()
def update_state(self, y_true, y_pred, sample_weight=None):
self.precision.update_state(y_true, y_pred, sample_weight)
self.recall.update_state(y_true, y_pred, sample_weight)
def result(self):
p, r = self.precision.result(), self.recall.result()
return 2 * p * r / (p + r + 1e-8)
def reset_state(self):
self.precision.reset_state(); self.recall.reset_state()Epochs, Batch Size & Steps
When using tf.data.Dataset, do not pass batch_size — the dataset is already batched. steps_per_epoch tells fit how many batches constitute one epoch when the dataset repeats indefinitely. validation_steps does the same for validation. Use initial_epoch to resume training from a checkpoint — fit will skip the first N epochs but the optimizer state (LR schedule, momentum) continues correctly.
model.fit(x_train, y_train, batch_size=32, epochs=10)
# with tf.data.Dataset
model.fit(train_dataset, epochs=10, validation_data=val_dataset)
# Dataset of unknown size: use steps_per_epoch
model.fit(
train_dataset, epochs=10, steps_per_epoch=500,
validation_data=val_dataset, validation_steps=50,
)
# initial epoch (resume training)
model.fit(train_dataset, initial_epoch=5, epochs=10)Class & Sample Weights
class_weight scales the loss for each class to counter imbalance — equivalent to oversampling without the extra compute. sample_weight gives per-example control, useful for hard-example mining or masking padded positions. When using a tf.data.Dataset, include sample_weight as the third tuple element so the dataset yields (x, y, w). Use sklearn.utils.class_weight to compute balanced weights automatically.
class_weight = {0: 1.0, 1: 5.0, 2: 1.0} # class 1 is rare
model.fit(x_train, y_train, class_weight=class_weight, epochs=10)
sample_weight = np.where(y_train == 1, 5.0, 1.0)
model.fit(x_train, y_train, sample_weight=sample_weight, epochs=10)
# with Dataset: yield (x, y, sample_weight) tuples
ds = tf.data.Dataset.from_tensor_slices((x, y, w))
model.fit(ds, epochs=10)
from sklearn.utils import class_weight
cw = class_weight.compute_class_weight('balanced', classes=np.unique(y), y=y)
class_weight = dict(enumerate(cw))Predict & Evaluate
predict returns a NumPy array of predictions for the whole input — use batch_size to control memory. evaluate returns the loss plus any metrics from compile. For very large inputs, iterate batches manually with predict_on_batch. Calling the model directly (model(x)) returns a tensor and respects training=, while predict always runs in inference mode and returns NumPy. Use model(x, training=True) for test-time augmentation.
preds = model.predict(x_test, batch_size=64)
print(preds.shape) # (N, num_classes)
pred = model.predict(x_test[:1])
loss, acc = model.evaluate(x_test, y_test, batch_size=64, verbose=2)
test_loss, test_acc, test_prec = model.evaluate(test_dataset)
for batch in test_dataset:
preds = model.predict_on_batch(batch[0])
preds = model(x_test, training=True) # callable form, test-time augOverride train_step
Overriding train_step lets you customize what happens in each fit batch while keeping fit's epoch loop, callbacks, and distribution. This is the bridge between high-level fit and a fully custom loop — great for GANs, contrastive learning, or any non-standard objective. Use self.compiled_loss and self.compiled_metrics so the loss/metrics from compile() still work. test_step overrides validation behavior the same way.
import tensorflow as tf
class CustomModel(tf.keras.Model):
def train_step(self, data):
x, y = data
with tf.GradientTape() as tape:
y_pred = self(x, training=True)
loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
grads = tape.gradient(loss, self.trainable_variables)
self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
self.compiled_metrics.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)
self.compiled_loss(y, y_pred)
self.compiled_metrics.update_state(y, y_pred)
return {m.name: m.result() for m in self.metrics}
model = CustomModel(inputs, outputs)
model.compile(optimizer='adam', loss='...', metrics=['accuracy'])
model.fit(dataset, epochs=10)Callbacks
EarlyStopping
EarlyStopping halts training when the monitored metric stops improving, saving time and preventing overfitting. patience=N waits N epochs of no improvement before stopping. restore_best_weights=True rolls back to the best epoch at the end — without it you keep the (possibly worse) final weights. mode='min' for losses, 'max' for accuracy. Always set a generous max epochs and let the callback decide.
import tensorflow as tf
early = tf.keras.callbacks.EarlyStopping(
monitor='val_loss', min_delta=0.001, patience=5,
mode='auto', restore_best_weights=True, verbose=1,
)
model.fit(train_ds, validation_data=val_ds, epochs=100, callbacks=[early])ModelCheckpoint
ModelCheckpoint saves the model during training so you can recover from crashes or pick the best epoch. save_best_only=True keeps only the best checkpoint (by monitor metric) — set monitor to your val metric. save_weights_only=True is smaller and architecture-agnostic but you need the model class to reload. Format placeholders like {epoch:02d} embed the epoch number in the filename. save_freq='epoch' is the default; an integer N saves every N batches.
import tensorflow as tf
ckpt = tf.keras.callbacks.ModelCheckpoint(
filepath='checkpoints/epoch-{epoch:02d}-val-{val_loss:.3f}.keras',
monitor='val_loss', save_best_only=True, save_weights_only=False,
mode='min', save_freq='epoch', verbose=1,
)
ckpt = tf.keras.callbacks.ModelCheckpoint(
'weights.{epoch:02d}.h5', save_weights_only=True, save_freq='epoch',
)
model.fit(train_ds, validation_data=val_ds, epochs=50, callbacks=[ckpt])ReduceLROnPlateau
ReduceLROnPlateau drops the LR when the monitored metric plateaus — useful when you don't know the right schedule in advance. factor is the multiplier (0.5 = halve), patience is how many epochs to wait, min_lr is the floor. cooldown pauses monitoring right after a reduction to let the new LR settle. Combine with EarlyStopping: reduce LR first, then stop if no further improvement. verbose=1 prints each reduction.
import tensorflow as tf
reduce = tf.keras.callbacks.ReduceLROnPlateau(
monitor='val_loss', factor=0.5, patience=3, min_lr=1e-7,
mode='min', min_delta=1e-4, cooldown=0, verbose=1,
)
model.fit(train_ds, validation_data=val_ds, epochs=100, callbacks=[reduce])TensorBoard Callback
The TensorBoard callback streams training metrics, weight histograms, the model graph, and profiling data to a log directory. histogram_freq=1 logs weight distributions per epoch — set to 0 to save space. profile_batch=2 runs the TF Profiler on batch 2 to identify performance bottlenecks (input pipeline, op latency). Launch TensorBoard pointing at log_dir to view charts. Use distinct subdirectories per experiment to compare runs.
import tensorflow as tf
tb = tf.keras.callbacks.TensorBoard(
log_dir='./logs/fit', histogram_freq=1, write_graph=True,
write_images=False, update_freq='epoch',
profile_batch=2, embeddings_freq=0,
)
model.fit(train_ds, validation_data=val_ds, epochs=10, callbacks=[tb])
# launch: tensorboard --logdir=./logsLearningRateScheduler
LearningRateScheduler takes a function of (epoch, lr) and applies the new LR at each epoch start — simple but only per-epoch. For per-step schedules prefer a LearningRateSchedule object passed to the optimizer (see optimizers section) — it's called every batch and works inside tf.function. The Callback form is fine for coarse schedules. For warmup+decay, a schedule object is cleaner.
import tensorflow as tf
def schedule(epoch, lr):
if epoch < 10:
return lr
return lr * tf.math.exp(-0.1).numpy()
scheduler = tf.keras.callbacks.LearningRateScheduler(schedule, verbose=1)
class StepScheduler(tf.keras.callbacks.Callback):
def on_train_batch_begin(self, batch, logs=None):
step = self.model.optimizer.iterations.numpy()
new_lr = 1e-3 * (0.5 ** (step // 1000))
self.model.optimizer.learning_rate.assign(new_lr)
scheduler = tf.keras.callbacks.LearningRateScheduler(
lambda epoch: 1e-3 * 0.9 ** epoch)Custom Callback
Callbacks hook into the training lifecycle: on_epoch_begin/end, on_train_batch_begin/end, on_test_*, on_predict_*. Access the model via self.model and the optimizer via self.model.optimizer. logs is a dict of current metrics. Custom callbacks are great for logging, LR manipulation, early termination (self.model.stop_training = True), or syncing with external systems (W&B, MLflow). Keep them cheap — they run on every batch.
import tensorflow as tf
class DebugCallback(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
lr = self.model.optimizer.learning_rate
if hasattr(lr, 'numpy'):
lr = lr.numpy()
print(f"Epoch {epoch}: lr={lr}, logs={logs}")
def on_train_batch_end(self, batch, logs=None):
if batch % 100 == 0:
print(f" batch {batch}: loss={logs.get('loss'):.4f}")
def on_train_end(self, logs=None):
print("Training finished!")
model.fit(train_ds, epochs=5, callbacks=[DebugCallback()])tf.data Pipeline
tf.data.Dataset
tf.data.Dataset is the streaming data API for TF — it pipelines data from disk/CPU to the GPU efficiently. from_tensor_slices zips tensors along axis 0. The canonical chain is shuffle -> batch -> prefetch: shuffle randomizes, batch groups, prefetch overlaps the next batch's prep with current training. Use AUTOTUNE to let TF pick prefetch buffer sizes automatically.
import tensorflow as tf
ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
def gen():
for i in range(len(x_train)):
yield x_train[i], y_train[i]
ds = tf.data.Dataset.from_generator(gen, output_signature=(
tf.TensorSpec(shape=(28, 28), dtype=tf.float32),
tf.TensorSpec(shape=(), dtype=tf.int32),
))
ds = ds.shuffle(10000).batch(32).prefetch(tf.data.AUTOTUNE)
for x, y in ds:
print(x.shape, y.shape)
breakMap, Filter & Cache
map applies a function to each element — augmentation, normalization, tokenization. num_parallel_calls=AUTOTUNE parallelizes across CPU cores. cache() stores the dataset in memory (or file) so subsequent epochs skip the source — huge speedup for small-to-medium data, but cache BEFORE random ops or you'll repeat the same augmentations every epoch. The canonical order is cache -> map(augment) -> shuffle -> batch -> prefetch.
import tensorflow as tf
def augment(image, label):
image = tf.image.random_flip_left_right(image)
image = tf.image.random_brightness(image, 0.2)
return image, label
ds = tf.data.Dataset.from_tensor_slices((images, labels))
ds = ds.map(augment, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.map(lambda x, y: (tf.cast(x, tf.float32) / 255.0, y))
ds = ds.filter(lambda x, y: y < 10)
ds = ds.cache() # in memory
ds = ds.cache('/tmp/data.cache') # to file
ds = ds.cache().shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)Batch, Shuffle & Repeat
shuffle uses a buffer of size N: it fills the buffer, samples one, refills — so a buffer >= dataset size gives true random order. shuffle -> repeat -> batch reshuffles every epoch. For training with steps_per_epoch, use repeat() so the dataset never runs out. drop_remainder=True drops the last short batch — important for BatchNorm and fixed-shape models. Seed makes shuffling reproducible across iterations only if reshuffle_each_iteration=True with the same seed.
import tensorflow as tf
ds = tf.data.Dataset.from_tensor_slices((x, y))
ds = ds.batch(32, drop_remainder=False)
ds = ds.shuffle(buffer_size=10000, seed=42, reshuffle_each_iteration=True)
ds = ds.repeat() # infinite
ds = ds.repeat(5) # 5 epochs
# recommended order: shuffle -> repeat -> batch
ds = ds.shuffle(10000).repeat().batch(32)Prefetch & Performance
prefetch(AUTOTUNE) is the single biggest win — it lets the CPU prepare batch N+1 while the GPU trains on batch N, so neither waits. num_parallel_calls=AUTOTUNE on map parallelizes CPU-bound preprocessing. The optimal pipeline order is: source -> cache -> shuffle -> map(augment) -> batch -> prefetch. Use tf.data.experimental.AUTOTUNE (same as tf.data.AUTOTUNE) and let TF tune buffer sizes. Always profile before optimizing — the bottleneck is often elsewhere.
import tensorflow as tf
ds = ds.prefetch(tf.data.AUTOTUNE)
ds = ds.map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE)
# full recommended pipeline
ds = (tf.data.Dataset.from_tensor_slices((x, y))
.cache()
.shuffle(10000)
.map(augment, num_parallel_calls=tf.data.AUTOTUNE)
.batch(32)
.prefetch(tf.data.AUTOTUNE))Reading Files
TextLineDataset streams lines from text files (logs, CSV rows). make_csv_dataset parses CSVs with automatic column type inference. image_dataset_from_directory is the Keras equivalent of PyTorch's ImageFolder — it infers labels from subdirectory names and returns (image_batch, label_batch). For more control over image loading use tf.io.read_file + tf.io.decode_image inside a map function. Always batch the dataset before passing to fit.
import tensorflow as tf
ds = tf.data.TextLineDataset(['file1.txt', 'file2.txt'])
ds = tf.data.experimental.make_csv_dataset(
'data.csv', batch_size=32, label_name='label',
num_epochs=1, shuffle_buffer_size=1000,
)
ds = tf.keras.utils.image_dataset_from_directory(
'data/train', image_size=(224, 224), batch_size=32,
shuffle=True, seed=42, validation_split=0.2, subset='training',
)TFRecord
TFRecord is TF's binary record format — the fastest way to feed large datasets from disk. Each record is a serialized tf.train.Example (a dict of features). Reading is sequential I/O, which is much faster than random reads on cloud storage (GCS/S3). For images, store raw bytes and decode in map(). Use tf.data.TFRecordDataset with compression='GZIP' for smaller files. Sharding (multiple .tfrecord-0000N-of-0000M files) enables parallel reading.
import tensorflow as tf
def serialize(x, y):
feature = {
'x': tf.train.Feature(float_list=tf.train.FloatList(value=x.flatten())),
'y': tf.train.Feature(int64_list=tf.train.Int64List(value=[y])),
}
return tf.train.Example(features=tf.train.Features(feature=feature)).SerializeToString()
with tf.io.TFRecordWriter('data.tfrecord') as w:
for x, y in dataset:
w.write(serialize(x, y))
def parse(serialized):
feature = {'x': tf.io.FixedLenFeature([784], tf.float32),
'y': tf.io.FixedLenFeature([], tf.int64)}
parsed = tf.io.parse_single_example(serialized, feature)
return parsed['x'], parsed['y']
ds = tf.data.TFRecordDataset('data.tfrecord').map(parse).batch(32)Save & Load
SavedModel Format
SavedModel is the production-grade format: it bundles the architecture, weights, and a traced graph so the model can be loaded in TF Serving, TFLite, TF.js, or C++ without the original Python code. Use tf.keras.models.load_model to reload into Python. Pass an input_signature to define the serving signature explicitly — required for TF Serving. SavedModel is a directory, not a single file.
import tensorflow as tf
model.save('saved_model') # SavedModel directory
loaded = tf.keras.models.load_model('saved_model')
loaded.evaluate(x_test, y_test)
# export for serving (explicit signature)
@tf.function(input_signature=[tf.TensorSpec([None, 784], tf.float32)])
def serve(x):
return {'outputs': model(x)}
tf.saved_model.save(model, 'serving_model', signatures=serve)Keras Format (.keras)
The .keras format (Keras 3 default) is a single zipped file containing the architecture config, weights, and optimizer state — easy to share and version. .h5 is the legacy HDF5 format, still supported but not recommended for new projects. SavedModel (a directory) is required for deployment (TFLite, TF Serving). For custom layers, .keras saves the config so the layer can be rebuilt — make sure get_config() is correct.
import tensorflow as tf
model.save('model.keras') # Keras 3 default
loaded = tf.keras.models.load_model('model.keras')
# legacy HDF5 format
model.save('model.h5')
loaded = tf.keras.models.load_model('model.h5')
# .keras — Keras 3 default, supports custom layers via config
# .h5 — legacy, weights + architecture but limited custom layer support
# SavedModel — directory, best for deployment/servingSave & Load Weights Only
save_weights stores only the parameter tensors — much smaller than a full model and decoupled from the architecture. You must rebuild the model with the same structure before loading. skip_mismatch=True (Keras 3) loads only matching layers, useful for transfer learning. The checkpoint format ('ckpt' prefix) creates checkpoint/weights-of-ckpt-1.index files — TF's native checkpoint format used internally by fit callbacks.
import tensorflow as tf
model.save_weights('weights.keras')
model = build_model() # must match the saved architecture
model.load_weights('weights.keras')
model.save_weights('weights.h5')
model.load_weights('weights.h5')
# load weights from one model into another (partial)
model.load_weights('weights.keras', skip_mismatch=True, by_name=True)
model.save_weights('ckpt')
model.load_weights('ckpt')Checkpoint Manager
Checkpoint + CheckpointManager give you full control over saving in custom training loops — keep the last N checkpoints, restore the latest automatically, and version by step. expect_partial() silences warnings about unrestored slots (e.g. optimizer momentum for new variables). Checkpoints store object graph, not class definitions — you still need the model code to restore. CheckpointManager.max_to_keep=3 auto-deletes old checkpoints.
import tensorflow as tf
ckpt = tf.train.Checkpoint(model=model, optimizer=optimizer, step=optimizer.iterations)
manager = tf.train.CheckpointManager(
ckpt, directory='./ckpts', max_to_keep=3,
checkpoint_name='step-{step}',
)
manager.save()
ckpt.restore(manager.latest_checkpoint).expect_partial()
# in a custom loop: save every N steps
if step % 1000 == 0:
manager.save()Load with Custom Objects
When loading a model that uses custom layers, losses, or metrics, Keras needs to know how to reconstruct them. Pass custom_objects dict mapping names to classes/functions. For cleaner code, decorate classes with @register_keras_serializable so they're registered globally and don't need custom_objects. Make sure get_config returns all constructor args so the object can be rebuilt exactly — Keras calls from_config(config) under the hood.
import tensorflow as tf
class MyLayer(tf.keras.layers.Layer):
...
custom_objects = {'MyLayer': MyLayer, 'focal_loss': FocalLoss}
loaded = tf.keras.models.load_model(
'model.keras', custom_objects=custom_objects,
)
# or register globally
@tf.keras.utils.register_keras_serializable(package='my_pkg')
class MyLayer(tf.keras.layers.Layer):
...Export for Serving
TF Serving expects versioned directories (serving/1/, serving/2/) so it can hot-swap models. The serving signature defines the input/output tensor names and shapes clients use. model.export() (Keras 3) wraps this automatically. For custom preprocessing, wrap the model in a tf.Module and define __call__ with an input_signature that includes the raw input (e.g. uint8 image) and returns the final output. Name tensors so clients can address them.
import tensorflow as tf
# Keras 3 export API (recommended)
model.export('serving/1')
class ServingModule(tf.Module):
def __init__(self, model):
self.model = model
@tf.function(input_signature=[tf.TensorSpec([None, 784], tf.float32, name='inputs')])
def __call__(self, x):
return {'outputs': self.model(x)}
module = ServingModule(model)
tf.saved_model.save(module, 'serving/1',
signatures={'serving_default': module.__call__})
loaded = tf.saved_model.load('serving/1')
print(list(loaded.signatures.keys())) # ['serving_default']CNN
Conv2D & Pooling
The canonical CNN block in Keras: Conv2D -> BatchNormalization -> ReLU -> MaxPooling. padding='same' keeps spatial size for stride 1. GlobalAveragePooling2D collapses spatial dims to a vector — more parameter-efficient than Flatten + Dense and less prone to overfitting. Keras uses NHWC (channels-last) by default, so Input shape is (H, W, C) without a batch dimension. Strided convs can replace pooling for learnable downsampling.
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(28, 28, 1)),
tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.MaxPooling2D(), # 28 -> 14
tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.MaxPooling2D(), # 14 -> 7
tf.keras.layers.Conv2D(128, 3, padding='same', activation='relu'),
tf.keras.layers.GlobalAveragePooling2D(), # (7,7,128) -> (128,)
tf.keras.layers.Dense(10, activation='softmax'),
])
model.summary()Building a CNN
A VGG-style CNN with doubled filters at each stage (32 -> 64 -> 128) and halved spatial size via MaxPooling. Two convs per stage (VGG block) extracts richer features than one. Rescaling inside the model makes it deployable — raw uint8 images go straight in. GlobalAveragePooling + Dropout regularizes the head. Compile with sparse_categorical_crossentropy because labels are integer class indices.
import tensorflow as tf
inputs = tf.keras.Input(shape=(32, 32, 3))
x = tf.keras.layers.Rescaling(1./255)(inputs)
for filters in [32, 64, 128]:
x = tf.keras.layers.Conv2D(filters, 3, padding='same', activation='relu')(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Conv2D(filters, 3, padding='same', activation='relu')(x)
x = tf.keras.layers.MaxPooling2D()(x)
x = tf.keras.layers.Dropout(0.25)(x)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dense(256, activation='relu')(x)
x = tf.keras.layers.Dropout(0.5)(x)
outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
model = tf.keras.Model(inputs, outputs)Feature Maps & Visualization
Create a Model with the same input but multiple outputs (the conv layer activations) to inspect what each layer learns. Early layers detect edges/colors, deeper layers detect textures/parts. This is invaluable for debugging — dead feature maps (all zeros) indicate broken ReLUs or learning rate issues. For Grad-CAM (class activation maps), use tf-keras-vis or implement gradient-weighted class activations to see which regions drove a prediction.
import tensorflow as tf
import matplotlib.pyplot as plt
feature_model = tf.keras.Model(
inputs=model.inputs,
outputs=[layer.output for layer in model.layers if 'conv2d' in layer.name],
)
img = x_test[0:1]
features = feature_model.predict(img)
print([f.shape for f in features])
fig, axes = plt.subplots(4, 8, figsize=(12, 6))
for i, ax in enumerate(axes.flat):
if i < features[0].shape[-1]:
ax.imshow(features[0][0, :, :, i], cmap='viridis')
ax.axis('off')Pretrained CNN Backbones
keras.applications exposes ResNet, EfficientNet, MobileNet, ViT and more with ImageNet weights. include_top=False drops the 1000-class head so you can attach your own. preprocess_input applies the exact normalization the backbone was trained with — call it inside the model for deployment. Freeze the backbone (trainable=False) and set training=False on the call so BatchNorm uses inference stats — critical for transfer learning with small datasets.
import tensorflow as tf
backbone = tf.keras.applications.ResNet50(
include_top=False, weights='imagenet',
input_shape=(224, 224, 3), pooling='avg',
)
backbone.trainable = False
inputs = tf.keras.Input(shape=(224, 224, 3))
x = tf.keras.applications.resnet50.preprocess_input(inputs)
x = backbone(x, training=False)
x = tf.keras.layers.Dense(256, activation='relu')(x)
x = tf.keras.layers.Dropout(0.5)(x)
outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
model = tf.keras.Model(inputs, outputs)Residual Connections
Residual (skip) connections add the input to the output of a block, letting gradients flow directly to earlier layers — this is what enables training very deep networks (ResNet). The shortcut must match the output's channels and spatial size; use a 1x1 conv with stride to adjust when they differ. Add() merges the two paths. Use functional API (not Sequential) for any architecture with skip connections. ResNet-50 stacks many such blocks.
import tensorflow as tf
def residual_block(x, filters, stride=1):
shortcut = x
x = tf.keras.layers.Conv2D(filters, 3, strides=stride, padding='same', use_bias=False)(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.ReLU()(x)
x = tf.keras.layers.Conv2D(filters, 3, padding='same', use_bias=False)(x)
x = tf.keras.layers.BatchNormalization()(x)
if stride != 1 or shortcut.shape[-1] != filters:
shortcut = tf.keras.layers.Conv2D(filters, 1, strides=stride, use_bias=False)(shortcut)
shortcut = tf.keras.layers.BatchNormalization()(shortcut)
x = tf.keras.layers.Add()([x, shortcut])
return tf.keras.layers.ReLU()(x)
inputs = tf.keras.Input(shape=(32, 32, 3))
x = residual_block(inputs, 64)
x = residual_block(x, 128, stride=2)
model = tf.keras.Model(inputs, x)Detection Head
Detection heads predict per-anchor or per-pixel boxes plus class probabilities. The functional API makes multi-output models easy: return a dict of named tensors. For real projects, use the TensorFlow Object Detection API or KerasCV (detection, segmentation) — they handle anchor matching, NMS, and evaluation. Writing a full detector from scratch is an advanced exercise: handle box encoding, IoU loss, and focal loss carefully.
import tensorflow as tf
# YOLO-style detection head: conv to (grid * grid * (5 + num_classes))
def detection_head(x, num_classes, num_anchors=3):
out = tf.keras.layers.Conv2D(
num_anchors * (5 + num_classes), 1, activation='linear')(x)
return out
# anchor-free alternative: predict (cx, cy, w, h, obj, class_logits)
def center_net_head(features, num_classes):
cls = tf.keras.layers.Conv2D(num_classes, 1, activation='sigmoid', name='cls')(features)
wh = tf.keras.layers.Conv2D(2, 1, name='wh')(features)
offset = tf.keras.layers.Conv2D(2, 1, name='offset')(features)
return {'cls': cls, 'wh': wh, 'offset': offset}RNN & LSTM
LSTM & GRU Layers
LSTM/GRU input shape is (batch, timesteps, features). return_sequences=False returns the last output (for classification); True returns the full sequence (for stacking or seq2seq). Bidirectional concatenates forward and backward outputs, doubling the feature size. stateful=True carries hidden state across batches — useful for very long sequences split into chunks, but requires fixed batch_size and manual state resets between epochs.
import tensorflow as tf
lstm = tf.keras.layers.LSTM(128, return_sequences=False)
x = tf.random.normal((4, 10, 64))
out = lstm(x) # (4, 128)
gru = tf.keras.layers.GRU(128)
out = gru(x) # (4, 128)
lstm = tf.keras.layers.LSTM(128, return_sequences=True)
out = lstm(x) # (4, 10, 128)
bi = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64))
out = bi(x) # (4, 128) — 64*2
lstm = tf.keras.layers.LSTM(64, stateful=True)Stacked RNNs
Stack RNNs by setting return_sequences=True on all but the last LSTM/GRU. mask_zero=True in the Embedding layer generates a mask that RNNs respect — they skip padded positions, so you can batch variable-length sequences without affecting the hidden state. Dropout between RNN layers regularizes; recurrent_dropout inside an RNN is slower (no cuDNN) — prefer input dropout. For text classification this biLSTM is a strong baseline.
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Embedding(10000, 128, input_length=100, mask_zero=True),
tf.keras.layers.LSTM(128, return_sequences=True),
tf.keras.layers.Dropout(0.3),
tf.keras.layers.LSTM(64),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax'),
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')Text Classification with RNN
A standard text classifier: Embedding -> biLSTM -> biLSTM -> Dense -> sigmoid. mask_zero=True is critical — without it the LSTM processes padding tokens, biasing the hidden state. Bidirectional doubles the representational power by reading the sequence both ways. For binary classification use sigmoid + binary_crossentropy. For long sequences (>500 tokens), transformers outperform LSTMs but need more data and compute.
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(200,)),
tf.keras.layers.Embedding(10000, 128, mask_zero=True),
tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64, return_sequences=True)),
tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(32)),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(1, activation='sigmoid'),
])
model.compile(optimizer=tf.keras.optimizers.Adam(1e-3),
loss='binary_crossentropy', metrics=['accuracy'])
model.fit(train_ds, validation_data=val_ds, epochs=5)Encoder-Decoder Seq2Seq
Encoder-decoder LSTM for seq2seq (translation, summarization). The encoder's final hidden state initializes the decoder, which generates the target sequence token by token. Train with teacher forcing: feed the shifted target sequence as decoder input. For inference, write a separate loop that feeds each predicted token back as the next input. Attention (Bahdanau/Luong) dramatically improves seq2seq and is the conceptual ancestor of transformers.
import tensorflow as tf
encoder_input = tf.keras.Input(shape=(None,), name='enc_in')
enc_emb = tf.keras.layers.Embedding(vocab_src, 128)(encoder_input)
encoder = tf.keras.layers.LSTM(256, return_state=True)
enc_out, state_h, state_c = encoder(enc_emb)
decoder_input = tf.keras.Input(shape=(None,), name='dec_in')
dec_emb = tf.keras.layers.Embedding(vocab_tgt, 128)(decoder_input)
decoder_lstm = tf.keras.layers.LSTM(256, return_sequences=True, return_state=True)
dec_out, _, _ = decoder_lstm(dec_emb, initial_state=[state_h, state_c])
decoder_dense = tf.keras.layers.Dense(vocab_tgt, activation='softmax')
outputs = decoder_dense(dec_out)
model = tf.keras.Model([encoder_input, decoder_input], outputs)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')Time Series Forecasting
For time series, build windowed (input, target) pairs: the last value in each window is the target, the rest are inputs. window(shift=1) creates overlapping windows; drop_remainder keeps them the same size. LSTM learns temporal patterns; a final Dense(1) predicts the next value. For multi-step forecasting, predict a horizon vector or use a seq2seq architecture. Always standardize the series using only training statistics.
import tensorflow as tf
def make_dataset(series, window=20, batch=32):
ds = tf.data.Dataset.from_tensor_slices(series)
ds = ds.window(window + 1, shift=1, drop_remainder=True)
ds = ds.flat_map(lambda w: w.batch(window + 1))
ds = ds.map(lambda w: (w[:-1][..., tf.newaxis], w[-1:]))
return ds.shuffle(1000).batch(batch).prefetch(tf.data.AUTOTUNE)
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(20, 1)),
tf.keras.layers.LSTM(32, return_sequences=False),
tf.keras.layers.Dense(1),
])
model.compile(optimizer='adam', loss='mse')
model.fit(make_dataset(train_series), epochs=10)Attention Mechanism
Attention lets the decoder focus on relevant encoder states instead of relying on a single fixed vector — a major improvement for long sequences. Bahdanau (additive) and Luong (multiplicative) are the two classic variants. Keras provides tf.keras.layers.AdditiveAttention and Attention as ready-made layers. Attention is the conceptual bridge to transformers: multi-head self-attention is just attention applied to the same sequence.
import tensorflow as tf
class BahdanauAttention(tf.keras.layers.Layer):
def __init__(self, units):
super().__init__()
self.W1 = tf.keras.layers.Dense(units)
self.W2 = tf.keras.layers.Dense(units)
self.V = tf.keras.layers.Dense(1)
def call(self, query, values):
q = tf.expand_dims(query, 1)
scores = self.V(tf.nn.tanh(self.W1(values) + self.W2(q)))
attn = tf.nn.softmax(scores, axis=1)
context = tf.reduce_sum(attn * values, axis=1)
return context, attn
attn = tf.keras.layers.AdditiveAttention()
context = attn([query, value])Transformer
MultiHeadAttention
MultiHeadAttention is the core transformer layer — it projects q/k/v into multiple heads, attends in parallel, and recombines. For self-attention pass the same tensor as query, key, and value. attention_mask can be a causal mask (lower-triangular) for decoding or a padding mask (True = ignore). return_attention_scores=True returns the attention matrix for visualization. The output dim equals the input dim so blocks can stack.
import tensorflow as tf
mha = tf.keras.layers.MultiHeadAttention(
num_heads=8, key_dim=64, dropout=0.1, use_bias=True)
x = tf.random.normal((4, 10, 512))
out, attn_weights = mha(query=x, value=x, key=x,
return_attention_scores=True)
print(out.shape) # (4, 10, 512)
print(attn_weights.shape) # (4, 8, 10, 10)
# causal mask (decoder self-attention)
mask = tf.linalg.band_part(tf.ones((10, 10)), -1, 0)
out = mha(query=x, value=x, key=x, attention_mask=mask)
# key padding mask (ignore pad positions)
padding = tf.cast(tokens == 0, tf.bool)[:, tf.newaxis, tf.newaxis, :]
out = mha(query=x, value=x, key=x, attention_mask=padding)Transformer Encoder Block
A transformer encoder block is: multi-head self-attention + add & norm + feed-forward (MLP) + add & norm. Pre-LN (norm before sublayer) is more stable for training from scratch; post-LN needs warmup. The FFN is a 2-layer MLP with GELU, expanding to 'ff' hidden units and back to 'dim'. Stack N of these for a full encoder. epsilon=1e-6 in LayerNorm avoids division by zero for small variances.
import tensorflow as tf
class TransformerBlock(tf.keras.layers.Layer):
def __init__(self, dim=512, heads=8, ff=2048, dropout=0.1, **kwargs):
super().__init__(**kwargs)
self.att = tf.keras.layers.MultiHeadAttention(num_heads=heads, key_dim=dim // heads)
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(ff, activation='gelu'),
tf.keras.layers.Dense(dim),
])
self.norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.drop1 = tf.keras.layers.Dropout(dropout)
self.drop2 = tf.keras.layers.Dropout(dropout)
def call(self, x, training=False, mask=None):
attn = self.att(x, x, x, attention_mask=mask)
x = self.norm1(x + self.drop1(attn, training=training))
ffn = self.ffn(x)
return self.norm2(x + self.drop2(ffn, training=training))Positional Encoding
Transformers are permutation-invariant — they need positional info to know token order. Sinusoidal encodings add a fixed pattern based on position and dimension; they generalize to longer sequences than seen in training. Add the positional encoding to the embedding output. Alternatively use learned position embeddings (an Embedding layer indexed by position) or relative position encodings (used in modern transformers like T5). For images (ViT), add a learnable position embedding per patch.
import tensorflow as tf
import numpy as np
class PositionalEncoding(tf.keras.layers.Layer):
def __init__(self, max_len=5000, dim=512, **kwargs):
super().__init__(**kwargs)
pos = np.arange(max_len)[:, np.newaxis]
i = np.arange(dim)[np.newaxis, :]
angle = pos / np.power(10000, (2 * (i // 2)) / np.float32(dim))
angle[:, 0::2] = np.sin(angle[:, 0::2])
angle[:, 1::2] = np.cos(angle[:, 1::2])
self.pos = tf.constant(angle[np.newaxis, ...], dtype=tf.float32)
def call(self, x):
return x + self.pos[:, :tf.shape(x)[1], :]
model = tf.keras.Sequential([
tf.keras.layers.Embedding(10000, 512),
PositionalEncoding(max_len=200, dim=512),
tf.keras.layers.Dropout(0.1),
])Text Classification Transformer
A from-scratch transformer text classifier: embedding + positional encoding + N encoder blocks + global average pool + Dense head. GlobalAveragePooling1D collapses the sequence dimension (use a CLS token for a learnable aggregation). mask_zero=True propagates a mask that MultiHeadAttention respects — pass it through manually if needed. For real NLP, use a pretrained transformer (BERT, T5) from HuggingFace — training from scratch needs huge data.
import tensorflow as tf
class TransformerEncoder(tf.keras.layers.Layer):
def __init__(self, dim, heads, ff, num_layers, dropout=0.1, **kwargs):
super().__init__(**kwargs)
self.blocks = [TransformerBlock(dim, heads, ff, dropout)
for _ in range(num_layers)]
def call(self, x, training=False, mask=None):
for block in self.blocks:
x = block(x, training=training, mask=mask)
return x
inputs = tf.keras.Input(shape=(200,))
x = tf.keras.layers.Embedding(10000, 128, mask_zero=True)(inputs)
x = PositionalEncoding(200, 128)(x)
x = TransformerEncoder(128, 4, 512, 4)(x)
x = tf.keras.layers.GlobalAveragePooling1D()(x)
x = tf.keras.layers.Dropout(0.3)(x)
outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)
model = tf.keras.Model(inputs, outputs)Vision Transformer (ViT)
A Vision Transformer patchifies the image with a strided Conv2d (kernel=stride=patch_size), treats patches as tokens, prepends a CLS token, adds positional embeddings, and runs transformer encoder blocks. The CLS token's final hidden state is the global image representation for classification. For real use, load a pretrained ViT from keras.applications or HuggingFace — they need ImageNet-21k pretraining to work well. Patch size 16 and dim 768 are ViT-Base defaults.
import tensorflow as tf
class ViTPatchEmbed(tf.keras.layers.Layer):
def __init__(self, patch_size=16, dim=768, **kwargs):
super().__init__(**kwargs)
self.proj = tf.keras.layers.Conv2D(dim, patch_size, strides=patch_size)
def call(self, x):
x = self.proj(x)
shape = tf.shape(x)
x = tf.reshape(x, [shape[0], -1, x.shape[-1]])
return x
inputs = tf.keras.Input(shape=(224, 224, 3))
x = ViTPatchEmbed(16, 768)(inputs)
x = x + tf.Variable(tf.random.normal((1, 196, 768)))
cls = tf.Variable(tf.random.normal((1, 1, 768)))
x = tf.keras.layers.Concatenate(axis=1)([cls, x])
for _ in range(6):
x = TransformerBlock(768, 12, 3072, 0.1)(x)
x = tf.keras.layers.Lambda(lambda t: t[:, 0])(x)
outputs = tf.keras.layers.Dense(10, activation='softmax')(x)
model = tf.keras.Model(inputs, outputs)Transformer Decoder
A decoder block has three sublayers: masked self-attention (causal mask so position t can't see the future), cross-attention (query from decoder, key/value from encoder output), and an FFN. Each sublayer has a residual + LayerNorm. Cross-attention is what lets the decoder 'look back' at the source sequence. For decoder-only models (GPT), drop the cross-attention sublayer. Stack N blocks for a full decoder. Inference needs an autoregressive loop feeding predicted tokens back.
import tensorflow as tf
class DecoderBlock(tf.keras.layers.Layer):
def __init__(self, dim, heads, ff, dropout=0.1, **kwargs):
super().__init__(**kwargs)
self.self_att = tf.keras.layers.MultiHeadAttention(heads, dim // heads)
self.cross_att = tf.keras.layers.MultiHeadAttention(heads, dim // heads)
self.ffn = tf.keras.Sequential([
tf.keras.layers.Dense(ff, activation='gelu'),
tf.keras.layers.Dense(dim),
])
self.norm1 = self.norm2 = self.norm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
self.drop = tf.keras.layers.Dropout(dropout)
def call(self, x, enc_out, training=False, causal_mask=None, padding_mask=None):
attn = self.self_att(x, x, x, attention_mask=causal_mask)
x = self.norm1(x + self.drop(attn, training=training))
cross = self.cross_att(query=x, value=enc_out, key=enc_out, attention_mask=padding_mask)
x = self.norm2(x + self.drop(cross, training=training))
x = self.norm3(x + self.drop(self.ffn(x), training=training))
return xTensorBoard
Setup & Callback
The TensorBoard callback streams metrics, graphs, and profiling data to log_dir during fit. histogram_freq=1 logs weight distributions per epoch — set to 0 to save space. profile_batch=2 runs the TF Profiler on that batch to identify bottlenecks (input pipeline, op latency, memory). Use distinct subdirectories per experiment (logs/exp1, logs/exp2) so TensorBoard overlays them — perfect for hyperparameter sweeps. Launch with `tensorboard --logdir=logs`.
import tensorflow as tf
tb = tf.keras.callbacks.TensorBoard(
log_dir='logs/experiment_1', histogram_freq=1,
write_graph=True, profile_batch=2,
)
model.fit(train_ds, validation_data=val_ds, epochs=10, callbacks=[tb])
# launch tensorboard
# tensorboard --logdir=logs
import datetime
log_dir = 'logs/fit/' + datetime.datetime.now().strftime('%Y%m%d-%H%M%S')Custom Scalars (Summary Writer)
tf.summary.create_file_writer gives you a low-level handle for logging custom scalars in custom training loops. Use writer.as_default() and tf.summary.scalar(name, value, step=step) to write. Step must be an integer (usually the optimizer's iteration counter). Flush periodically so events reach disk. Use namespaces (train/loss, val/loss) to organize charts in the TensorBoard UI. Log LR, gradient norms, and any custom metric you care about.
import tensorflow as tf
writer = tf.summary.create_file_writer('logs/custom')
with writer.as_default():
tf.summary.scalar('train/loss', loss, step=step)
tf.summary.scalar('train/lr', lr, step=step)
tf.summary.scalar('val/accuracy', acc, step=step)
tf.summary.text('config', 'lr=1e-3, batch=32', step=0)
writer.flush()
# in a custom training loop
for step, (x, y) in enumerate(train_ds):
# ... train step ...
if step % 100 == 0:
with writer.as_default():
tf.summary.scalar('loss', loss, step=step)Images & Histograms
tf.summary.image logs images for visual QA — pass shape (B, H, W, C) with B up to max_outputs. Histograms show weight and gradient distributions over time; sudden spikes hint at instability. Log a few validation samples per epoch to see qualitative progress, not just scalar metrics. Audio is useful for speech models. All summaries take a step argument that becomes the x-axis in TensorBoard. Keep logging cheap by sampling every N steps.
import tensorflow as tf
writer = tf.summary.create_file_writer('logs/viz')
with writer.as_default():
tf.summary.image('sample', img[None, ...], step=step)
tf.summary.image('batch', images[:8], step=step, max_outputs=8)
for var in model.trainable_variables:
tf.summary.histogram('weights/' + var.name, var, step=step)
tf.summary.audio('waveform', audio[None, ...], sample_rate=16000, step=step)
writer.flush()Hyperparameter Tuning
The HParams plugin organizes hyperparameter sweeps in TensorBoard — a table view lets you filter runs and find the best config. Define HParam objects with ranges/discrete values, log a config once with hp.hparams_config, then for each run log hp.hparams(hparams) plus the final metric. For serious sweeps use KerasTuner (RandomSearch, Hyperband, BayesianOptimization) which automates the search and integrates with TensorBoard.
import tensorflow as tf
# from tensorboard.plugins.hparams import api as hp
# (run: pip install tensorboard-plugin-hparams)
HP_LR = hp.HParam('lr', hp.RealInterval(1e-4, 1e-2))
HP_DROPOUT = hp.HParam('dropout', hp.Discrete([0.1, 0.3, 0.5]))
METRIC_ACC = 'accuracy'
with tf.summary.create_file_writer('logs/hparam_tuning').as_default():
hp.hparams_config(hparams=[HP_LR, HP_DROPOUT], metrics=[hp.Metric(METRIC_ACC)])
for lr in [1e-3, 5e-3]:
for dropout in [0.1, 0.3]:
hparams = {HP_LR: lr, HP_DROPOUT: dropout}
run_name = f"lr-{lr}-drop-{dropout}"
with tf.summary.create_file_writer('logs/hparam_tuning/' + run_name).as_default():
hp.hparams(hparams)
acc = train(hparams)
tf.summary.scalar(METRIC_ACC, acc, step=1)Profiling Performance
The TF Profiler shows where time is spent: per-op, per-device, and in the input pipeline. profile_batch='10,20' profiles batches 10-20 (skip the first few so warmup is done). The Overview page gives actionable recommendations (e.g. 'input pipeline is the bottleneck'). The Memory Viewer shows peak memory and which tensors dominate. The Trace Viewer is a Chrome-style timeline for deep dives. Always profile on the target hardware — CPU vs GPU bottlenecks differ.
import tensorflow as tf
# profile via callback
tb = tf.keras.callbacks.TensorBoard(
log_dir='logs/profile', profile_batch='10,20',
)
# or use the Profiler API directly
tf.profiler.experimental.start('logs/profile')
# ... run a few batches ...
tf.profiler.experimental.stop()
# in TensorBoard, open the Profile tab to see:
# - op-level time breakdown
# - input pipeline analysis (tf.data bottleneck)
# - memory viewer
# - overview page with recommendationsGPU / TPU Training
GPU Detection & Memory
By default TF grabs all GPU memory at startup, blocking other processes. set_memory_growth(True) allocates only what's needed — essential for sharing a GPU. set_logical_device_configuration limits memory or splits one GPU into multiple logical devices (useful for testing multi-GPU code on one GPU). tf.device('/GPU:1') places ops on a specific device. CPU ops run on '/CPU:0'. Mismatched device placement causes implicit copies that hurt performance.
import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
print('GPUs:', gpus)
# set memory growth (don't grab all GPU memory at start)
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
# or limit GPU memory explicitly
tf.config.set_logical_device_configuration(
gpus[0],
[tf.config.LogicalDeviceConfiguration(memory_limit=4096)],
)
with tf.device('/GPU:1'):
x = tf.random.normal((1000, 1000))
y = tf.matmul(x, x)Distribution Strategy
Distribution strategies abstract away the device topology. MirroredStrategy is for single-machine multi-GPU (mirrors variables across GPUs, all-reduces gradients) — the easiest multi-GPU path. MultiWorkerMirroredStrategy extends to multiple machines. TPUStrategy runs on Google TPUs (Colab/Cloud). Always create the model inside strategy.scope() so variables are mirrored correctly. The same fit/predict code works across strategies.
import tensorflow as tf
strategy = tf.distribute.OneDeviceStrategy(device='/gpu:0')
strategy = tf.distribute.MirroredStrategy()
print('num replicas:', strategy.num_replicas_in_sync)
# 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)
with strategy.scope():
model = build_model()
model.compile(optimizer='adam', loss='...', metrics=['accuracy'])
model.fit(train_ds, epochs=10)Multi-GPU Training
MirroredStrategy mirrors model variables on each GPU and splits each batch across replicas — the effective batch size is global_batch (e.g. 64 per GPU * 4 GPUs = 256). Pass the global batch to the dataset; TF distributes shards automatically. fit() handles gradient all-reduce under the hood. For custom training loops, use strategy.experimental_distribute_dataset and strategy.run(train_step).
import tensorflow as tf
strategy = tf.distribute.MirroredStrategy()
print(f'using {strategy.num_replicas_in_sync} GPUs')
global_batch = 64 * strategy.num_replicas_in_sync
train_ds = make_dataset().batch(global_batch)
with strategy.scope():
model = build_model()
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
model.fit(train_ds, epochs=10, validation_data=val_ds)TPU Training (Colab)
TPUs are available in Google Colab (Change runtime type -> TPU) and Google Cloud. Initialize the TPU system once, then wrap model creation in strategy.scope(). TPUs are fastest with large batch sizes (1024+) and fixed shapes — use drop_remainder=True in batching. TPUs can't read local files directly; load data into memory or stream from GCS via tf.data. Mixed precision (bfloat16) gives a big speedup on TPU and needs no loss scaling.
import tensorflow as tf
try:
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)
print('TPUs:', strategy.num_replicas_in_sync)
except ValueError:
print('No TPU found')
strategy = tf.distribute.get_strategy()
with strategy.scope():
model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
tf.keras.layers.Dense(10, activation='softmax'),
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
model.fit(x_train, y_train, batch_size=1024, epochs=5)Mixed Precision
Mixed precision computes in float16 but keeps float32 master copies, giving 2-3x throughput on Tensor Core GPUs (Volta+) and halving memory. 'mixed_float16' needs loss scaling (Keras handles it automatically via LossScaleOptimizer). 'mixed_bfloat16' has the same range as float32 so no scaling is needed, but only Ampere+ GPUs and TPUs support it. Force the final layer to float32 for numerical stability (softmax/logits). Mixed precision is the easiest single performance win.
import tensorflow as tf
tf.keras.mixed_precision.set_global_policy('mixed_float16')
# or bfloat16 on Ampere+ / TPU (no loss scaling needed)
tf.keras.mixed_precision.set_global_policy('mixed_bfloat16')
model = tf.keras.Sequential([
tf.keras.layers.Dense(512, activation='relu', input_shape=(784,)),
tf.keras.layers.Dense(10, dtype='float32'),
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
print(model.layers[0].dtype_policy) # <Policy "mixed_float16">Data Augmentation
Keras Preprocessing Layers
Keras preprocessing layers (RandomFlip, RandomRotation, etc.) run on GPU and inside tf.function — far faster than PIL/numpy augmentation. Putting them inside the model means they're disabled at inference automatically (training=False). Alternatively, apply them in a tf.data map for more control. These layers are stateless and JIT-compile, so they don't slow inference. For stateful preprocessing (e.g. fitted scalers) use the Adapt API.
import tensorflow as tf
augment = tf.keras.Sequential([
tf.keras.layers.RandomFlip('horizontal'),
tf.keras.layers.RandomRotation(0.1),
tf.keras.layers.RandomZoom(0.1),
tf.keras.layers.RandomContrast(0.2),
tf.keras.layers.RandomTranslation(0.1, 0.1),
])
# put augmentation INSIDE the model (training-only, inactive at inference)
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(224, 224, 3)),
tf.keras.layers.Rescaling(1./255),
augment,
tf.keras.layers.Conv2D(32, 3, activation='relu'),
])
# or apply via dataset.map
ds = ds.map(lambda x, y: (augment(x, training=True), y))ImageDataGenerator (Legacy)
ImageDataGenerator is the legacy augmentation API — it works but runs on CPU and is slower than Keras preprocessing layers. Use it only when you need its directory-iteration convenience or legacy code. For new projects prefer keras.utils.image_dataset_from_directory + Keras preprocessing layers (RandomFlip etc.) — they're GPU-accelerated and integrate with tf.data. flow_from_directory infers class labels from subdirectory names like ImageFolder.
from tensorflow.keras.preprocessing.image import ImageDataGenerator
train_datagen = ImageDataGenerator(
rescale=1./255, rotation_range=20,
width_shift_range=0.2, height_shift_range=0.2,
shear_range=0.2, zoom_range=0.2,
horizontal_flip=True, fill_mode='nearest',
validation_split=0.2,
)
train_gen = train_datagen.flow_from_directory(
'data/train', target_size=(224, 224), batch_size=32,
class_mode='binary', subset='training',
)
model.fit(train_gen, epochs=10)Image Augmentation Ops
tf.image.* ops give fine-grained control for custom augmentation. random_brightness/contrast/saturation/hue change color without geometry. Cutout and Mixup are advanced augmentations that improve regularization; implement them with tf.image ops + tf.tensor_scatter_nd_update. Apply augmentation in a tf.data map with num_parallel_calls=AUTOTUNE so it runs on CPU workers and never blocks the GPU. Always clip back to valid pixel ranges after color jitter.
import tensorflow as tf
def augment(image, label):
image = tf.image.random_flip_left_right(image)
image = tf.image.random_brightness(image, max_delta=0.2)
image = tf.image.random_contrast(image, lower=0.8, upper=1.2)
image = tf.image.random_saturation(image, lower=0.8, upper=1.2)
image = tf.image.random_hue(image, max_delta=0.1)
image = tf.clip_by_value(image, 0.0, 1.0)
return image, label
ds = ds.map(augment, num_parallel_calls=tf.data.AUTOTUNE)Mixup & CutMix
Mixup linearly blends two images and their labels, creating virtual examples between classes — strong regularization that reduces overfitting and improves calibration. CutMix pastes a rectangular patch from one image onto another, preserving local structure better than Mixup. Both need the loss to accept soft (mixed) labels — use categorical_crossentropy, not sparse. Apply per batch after batching. Combine with standard geometric augmentation for best results.
import tensorflow as tf
def mixup(images, labels, alpha=0.2):
lam = tf.random.beta([], alpha, alpha)
batch_size = tf.shape(images)[0]
idx = tf.random.shuffle(tf.range(batch_size))
mixed_images = lam * images + (1 - lam) * tf.gather(images, idx)
mixed_labels = lam * labels + (1 - lam) * tf.gather(labels, idx)
return mixed_images, mixed_labels
ds = ds.batch(32).map(mixup, num_parallel_calls=tf.data.AUTOTUNE)Text & Tabular Augmentation
For text, simple word-level augmentations (synonym swap, random deletion/swap) are cheap and effective; for serious NLP use back-translation or models like EDA/AugBERT. For tabular data, mixup blends rows; gaussian noise on numeric features regularizes. Always apply augmentation only to training data, never validation/test. Keep augmentations label-preserving (or label-mixing like mixup) — corrupting labels hurts more than it helps.
import tensorflow as tf
import random
def text_augment(text):
words = text.split()
if random.random() < 0.2:
words = [w for w in words if random.random() > 0.2]
if len(words) > 1 and random.random() < 0.2:
i, j = random.sample(range(len(words)), 2)
words[i], words[j] = words[j], words[i]
return ' '.join(words)
def tabular_mixup(x, y, alpha=0.2):
lam = tf.random.beta([], alpha, alpha)
idx = tf.random.shuffle(tf.range(tf.shape(x)[0]))
return lam * x + (1 - lam) * tf.gather(x, idx), lam * y + (1 - lam) * tf.gather(y, idx)
def add_noise(x, std=0.01):
return x + tf.random.normal(tf.shape(x), stddev=std)Transfer Learning
Pretrained Models
keras.applications exposes ResNet, EfficientNet, MobileNet, ViT, and more with ImageNet weights. include_top=False drops the 1000-class head so you can attach your own. preprocess_input applies the exact normalization the backbone expects — call it inside the model for deployment. Freeze the backbone (trainable=False) and call base(x, training=False) so BatchNorm uses inference stats — critical for small datasets where BN statistics would otherwise be noisy.
import tensorflow as tf
base = tf.keras.applications.ResNet50(
include_top=False, weights='imagenet',
input_shape=(224, 224, 3), pooling='avg',
)
base.trainable = False
inputs = tf.keras.Input(shape=(224, 224, 3))
x = tf.keras.applications.resnet50.preprocess_input(inputs)
x = base(x, training=False)
x = tf.keras.layers.Dense(256, activation='relu')(x)
x = tf.keras.layers.Dropout(0.5)(x)
outputs = tf.keras.layers.Dense(num_classes, activation='softmax')(x)
model = tf.keras.Model(inputs, outputs)Feature Extraction
Feature extraction freezes the entire backbone and trains only a new head — the fastest transfer-learning baseline. Because the backbone is fixed, you can precompute features once (base.predict(train_ds)) and train a tiny classifier on them, saving huge compute. Set training=False when calling the backbone so BatchNorm uses running stats — otherwise small batches corrupt them. Good for small datasets (hundreds of images) where fine-tuning would overfit.
import tensorflow as tf
base = tf.keras.applications.MobileNetV2(
include_top=False, weights='imagenet', pooling='avg')
base.trainable = False
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(224, 224, 3)),
tf.keras.applications.mobilenet_v2.preprocess_input,
base,
tf.keras.layers.Dropout(0.3),
tf.keras.layers.Dense(num_classes, activation='softmax'),
])
model.compile(optimizer=tf.keras.optimizers.Adam(1e-3),
loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(train_ds, validation_data=val_ds, epochs=5)Fine-tuning
Fine-tuning unfreezes some backbone layers and trains them with a small LR so pretrained features adapt to your task. Always do feature extraction first (head-only) for a few epochs, THEN unfreeze — otherwise the random head's large gradients wreck the pretrained features. Use a much smaller LR (1e-5 vs 1e-3 for the head). Unfreeze from the top down (later layers are more task-specific; earlier layers are general). Keep BatchNorm layers frozen to preserve their running stats.
import tensorflow as tf
base = tf.keras.applications.ResNet50(include_top=False, weights='imagenet', pooling='avg')
base.trainable = True
# freeze everything except the last few blocks
for layer in base.layers[:-20]:
layer.trainable = False
model.compile(
optimizer=tf.keras.optimizers.Adam(1e-5), # 10x smaller than head
loss='sparse_categorical_crossentropy', metrics=['accuracy'],
)
model.fit(train_ds, validation_data=val_ds, epochs=5)Freezing & Unfreezing
Set layer.trainable=False to freeze it. After changing trainable, you MUST recompile the model so the optimizer picks up the new trainable set. Keep BatchNorm layers frozen during fine-tuning — their running stats are precious and small fine-tuning batches would corrupt them. Inspect trainable_weights vs weights to confirm what will update. Log the trainable parameter count so you don't accidentally freeze or unfreeze everything (a common silent bug).
import tensorflow as tf
base = tf.keras.applications.EfficientNetB0(
include_top=False, weights='imagenet', pooling='avg')
base.trainable = False
for layer in base.layers:
layer.trainable = 'block7' in layer.name
# freeze BatchNorm specifically (recommended during fine-tuning)
for layer in base.layers:
if isinstance(layer, tf.keras.layers.BatchNormalization):
layer.trainable = False
n_train = sum(tf.keras.backend.count_params(w) for w in model.trainable_weights)
n_total = sum(tf.keras.backend.count_params(w) for w in model.weights)
print(f"trainable: {n_train:,} / {n_total:,}")Multi-Task Transfer
Multi-task learning shares a backbone across related tasks, regularizing the representation and saving compute. Each head is a small branch from the shared features. Pass per-output losses and loss_weights in compile to balance task gradients (tune the weights so no task dominates). At fit time, pass targets as a dict matching output names. Multi-task transfer is powerful when you have one task with lots of data and another with little — the shared backbone benefits both.
import tensorflow as tf
base = tf.keras.applications.ResNet50(include_top=False, weights='imagenet')
base.trainable = False
inputs = tf.keras.Input(shape=(224, 224, 3))
x = tf.keras.applications.resnet50.preprocess_input(inputs)
features = base(x, training=False)
features = tf.keras.layers.GlobalAveragePooling2D()(features)
cls_out = tf.keras.layers.Dense(num_classes, activation='softmax', name='cls')(features)
box_out = tf.keras.layers.Dense(4, name='box')(features)
attr_out = tf.keras.layers.Dense(num_attr, activation='sigmoid', name='attr')(features)
model = tf.keras.Model(inputs, [cls_out, box_out, attr_out])
model.compile(
optimizer='adam',
loss={'cls': 'sparse_categorical_crossentropy', 'box': 'mse',
'attr': 'binary_crossentropy'},
loss_weights={'cls': 1.0, 'box': 5.0, 'attr': 0.5},
)Model Deployment
TFLite Converter
TFLite converts TF models to a flatbuffer format for mobile, embedded, and edge devices. Optimize.DEFAULT applies post-training quantization automatically; full int8 quantization needs a representative dataset to calibrate activation ranges and produces a model that runs on Edge TPU / microcontrollers. float16 keeps GPU compatibility and is safer if int8 hurts accuracy. The converter infers input shapes from the model; for dynamic shapes, fix them with concrete functions.
import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
with open('model.tflite', 'wb') as f:
f.write(tflite_model)
# with quantization (8-bit, ~4x smaller, faster on CPU)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# float16 quantization (GPU-friendly, 2x smaller)
converter.target_spec.supported_types = [tf.float16]
# full integer quantization (needs representative dataset)
converter.representative_dataset = rep_data
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
tflite_quant = converter.convert()TFLite Interpreter
The TFLite Interpreter runs .tflite models on CPU, GPU, Edge TPU, or microcontrollers. allocate_tensors prepares memory based on the fixed input shape — resize with resize_tensor_input for dynamic shapes. set_tensor / invoke / get_tensor is the inference loop. For batched inference, the model must be converted with a fixed batch size or use resize_tensor_input. TFLite has Python, C++, Java, Swift, and JavaScript bindings — the same .tflite file runs everywhere.
import tensorflow as tf
import numpy as np
interpreter = tf.lite.Interpreter(model_path='model.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(input_details)
# [{'name': 'input', 'shape': (1, 224, 224, 3), 'dtype': tf.float32, ...}]
input_data = np.expand_dims(img, axis=0).astype(np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output = interpreter.get_tensor(output_details[0]['index'])
print(output.shape, output.argmax())TF Serving
TF Serving hosts SavedModels behind a gRPC/REST API with versioning, batching, and GPU support. The directory structure serving/1/, serving/2/ enables hot model swaps. The Docker image is the easiest local deployment; for production use Vertex AI, SageMaker, or KServe. TF Serving automatically batches requests to maximize GPU utilization. The REST API uses {instances: [...]} for inputs and returns {predictions: [...]}; for named inputs/outputs use the signature_name parameter.
# save model in SavedModel format with versioning
# model.export('serving/1') # Keras 3
tf.saved_model.save(model, 'serving/1')
# run TF Serving via Docker
# docker run -p 8501:8501 --name tfserving \
# -v $(pwd)/serving:/models/my_model \
# -e MODEL_NAME=my_model \
# tensorflow/serving
# query via REST
import json, requests
data = json.dumps({'instances': x_test[:3].tolist()})
resp = requests.post(
'http://localhost:8501/v1/models/my_model:predict',
data=data, headers={'content-type': 'application/json'})
print(resp.json()['predictions'])Quantization-Aware Training
Quantization-Aware Training (QAT) inserts fake-quantization ops during training so the model adapts to the rounding noise it will see at int8 inference — usually recovering most of the accuracy lost by post-training quantization. Use tfmot (tensorflow-model-optimization) to wrap a trained model, fine-tune for a few epochs, then convert to TFLite. QAT is essential when int8 post-training quantization drops accuracy below your threshold (e.g. EfficientNet, object detectors). The resulting TFLite model runs on Edge TPU.
import tensorflow as tf
# (requires: pip install tensorflow-model-optimization)
import tensorflow_model_optimization as tfmot
model = tf.keras.applications.MobileNetV2(weights='imagenet')
annot_model = tfmot.quantization.keras.quantize_model(model)
annot_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
annot_model.fit(train_ds, epochs=3)
converter = tf.lite.TFLiteConverter.from_keras_model(annot_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_qat = converter.convert()TF.js & Edge Deployment
TF.js runs models in the browser or Node.js. The converter transforms SavedModel/.keras/.h5 into a model.json + binary weights directory. Layers models support the Keras API (fit, predict, even on-the-fly training); graph models are more compatible but less flexible. Browser inference uses WebGL/WASM; Node.js can use CUDA. TF.js is great for client-side ML (no server round-trip, privacy, offline). For mobile, prefer TFLite (smaller, faster, hardware acceleration).
# convert to TF.js format
# pip install tensorflowjs
# tensorflowjs_converter --input_format=keras model.keras tfjs_model/
# in JavaScript:
# import * as tf from '@tensorflow/tfjs';
# const model = await tf.loadLayersModel('tfjs_model/model.json');
# const out = model.predict(tf.tensor4d(img, [1, 224, 224, 3]));
# OR convert SavedModel -> TF.js
# tensorflowjs_converter \
# --input_format=tf_saved_model \
# --output_format=tfjs_graph_model \
# saved_model/ tfjs_model/
# TF.js layers model: full Keras-like API in the browser
# TF.js graph model: lower-level, supports more opsInference Optimization
For low-latency deployment, stack optimizations: pruning (zero out small weights, compress), knowledge distillation (train a small student to match a large teacher's logits), quantization (int8/float16), and XLA compilation (fuse ops). Each step has tradeoffs — always benchmark on the target hardware. Distillation usually preserves accuracy best; pruning+quantization give the biggest size reductions. XLA (jit_compile=True) fuses ops for a free 1.2-2x speedup on GPU/TPU.
import tensorflow as tf
# knowledge distillation: train a small student to mimic a large teacher
def distill_loss(teacher_logits, student_logits, temperature=3.0):
soft_targets = tf.nn.softmax(teacher_logits / temperature)
soft_probs = tf.nn.log_softmax(student_logits / temperature)
return tf.reduce_mean(
tf.keras.losses.categorical_crossentropy(soft_targets, soft_probs, from_logits=True))
# XLA compilation (faster inference)
@tf.function(jit_compile=True)
def predict(x):
return model(x)
# stack optimizations: prune + quantize + distill + XLA
# always benchmark latency and accuracy before/after each stepCustom Training Loop
Basic Custom Loop
A custom loop gives you full control: gradient manipulation, multiple models, custom objectives. The pattern is GradientTape -> loss -> grads -> apply_gradients, wrapped in @tf.function for speed. Don't forget model.losses — these are the regularization losses added by kernel_regularizer and should be added to the main loss. Track metrics with stateful Metric objects and reset_state() at each epoch. This is more verbose than fit but unlocks anything.
import tensorflow as tf
optimizer = tf.keras.optimizers.Adam(1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_acc = tf.keras.metrics.SparseCategoricalAccuracy(name='train_acc')
@tf.function
def train_step(x, y):
with tf.GradientTape() as tape:
logits = model(x, training=True)
loss = loss_fn(y, logits)
loss += sum(model.losses)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
train_loss.update_state(loss)
train_acc.update_state(y, logits)
return loss
for epoch in range(epochs):
for x, y in train_ds:
train_step(x, y)
print(f"epoch {epoch}: loss={train_loss.result():.4f} acc={train_acc.result():.4f}")
train_loss.reset_state(); train_acc.reset_state()Validation in Custom Loop
Validation mirrors training but with training=False (so dropout is off and BatchNorm uses running stats) and no gradient tape or optimizer. Use separate Metric objects for val so train and val don't mix. Wrap test_step in @tf.function for speed. Reset all metrics after logging each epoch. To save the best model, compare val_acc.result() to a running best and call model.save() when it improves.
import tensorflow as tf
val_loss = tf.keras.metrics.Mean(name='val_loss')
val_acc = tf.keras.metrics.SparseCategoricalAccuracy(name='val_acc')
@tf.function
def test_step(x, y):
logits = model(x, training=False)
loss = loss_fn(y, logits)
val_loss.update_state(loss)
val_acc.update_state(y, logits)
for epoch in range(epochs):
for x, y in train_ds:
train_step(x, y)
for x, y in val_ds:
test_step(x, y)
print(f"epoch {epoch}: val_loss={val_loss.result():.4f} val_acc={val_acc.result():.4f}")
train_loss.reset_state(); train_acc.reset_state()
val_loss.reset_state(); val_acc.reset_state()Distributed Custom Loop
A distributed custom loop uses strategy.run to run the per-replica step on each GPU and strategy.reduce to aggregate the loss. Divide the per-replica loss by num_replicas_in_sync so the summed gradient matches the global-batch gradient (otherwise the effective LR is scaled by the number of GPUs). Apply gradients normally — the strategy handles all-reduce under the hood. The same pattern works for MultiWorkerMirroredStrategy and TPUStrategy.
import tensorflow as tf
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
model = build_model()
optimizer = tf.keras.optimizers.Adam(1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
train_ds_dist = strategy.experimental_distribute_dataset(train_ds)
@tf.function
def distributed_train_step(x, y):
def step_fn(x, y):
with tf.GradientTape() as tape:
logits = model(x, training=True)
loss = loss_fn(y, logits)
loss += sum(model.losses)
loss /= strategy.num_replicas_in_sync
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
per_replica_loss = strategy.run(step_fn, args=(x, y))
return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None)
for epoch in range(epochs):
for x, y in train_ds_dist:
distributed_train_step(x, y)GAN Training Loop
GANs train two models adversarially: the generator tries to fool the discriminator, the discriminator tries to detect fakes. Use two GradientTapes (one per model) and two optimizers. beta_1=0.5 (lower than default) stabilizes GAN training. Train the discriminator and generator each step (or alternate ratios). GANs are notoriously unstable — monitor d_loss and g_loss, use spectral normalization, and consider Wasserstein GAN with gradient penalty if mode collapse occurs.
import tensorflow as tf
generator = build_generator()
discriminator = build_discriminator()
g_opt = tf.keras.optimizers.Adam(1e-4, beta_1=0.5)
d_opt = tf.keras.optimizers.Adam(1e-4, beta_1=0.5)
bce = tf.keras.losses.BinaryCrossentropy(from_logits=True)
@tf.function
def train_gan_step(real_images):
batch = tf.shape(real_images)[0]
noise = tf.random.normal((batch, 100))
with tf.GradientTape() as g_tape, tf.GradientTape() as d_tape:
fake = generator(noise, training=True)
real_logits = discriminator(real_images, training=True)
fake_logits = discriminator(fake, training=True)
d_loss = bce(tf.ones_like(real_logits), real_logits) + \
bce(tf.zeros_like(fake_logits), fake_logits)
g_loss = bce(tf.ones_like(fake_logits), fake_logits)
g_grads = g_tape.gradient(g_loss, generator.trainable_variables)
d_grads = d_tape.gradient(d_loss, discriminator.trainable_variables)
g_opt.apply_gradients(zip(g_grads, generator.trainable_variables))
d_opt.apply_gradients(zip(d_grads, discriminator.trainable_variables))
return g_loss, d_lossGradient Accumulation
True gradient accumulation in TF is tricky because GradientTape doesn't support summing grads across calls easily — you'd accumulate grads in tf.Variables manually. The simple approximation above scales the loss and applies gradients every step, which gives a similar effect for SGD but is not identical for Adam (the optimizer sees smaller, noisier gradients). For exact accumulation, accumulate gradients in persistent Variables and apply every N steps. This is mainly useful when GPU memory limits batch size.
import tensorflow as tf
accum_steps = 4
optimizer = tf.keras.optimizers.Adam(1e-3)
@tf.function
def accumulated_train_step(x, y):
with tf.GradientTape() as tape:
logits = model(x, training=True)
loss = loss_fn(y, logits) / accum_steps
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
# simple accumulation via averaging
for batch_idx, (x, y) in enumerate(train_ds):
accumulated_train_step(x, y)
# equivalent effective batch = accum_steps * batch_size
# (Adam moments are updated every step, so this is an approximation)Performance Optimization
tf.function
@tf.function traces a Python function into a static graph — faster on GPU and required for distribution. The first call per input signature traces (slow); later calls reuse the graph. Avoid Python control flow that depends on tensor VALUES (if tensor > 0: ...) inside tf.function — it forces retracing. Use tf.cond / tf.while_loop or Python control flow that depends only on shapes/Python literals. input_signature pins shapes to prevent retracing entirely. reduce_retracing=True broadens the cache.
import tensorflow as tf
@tf.function
def train_step(x, y):
with tf.GradientTape() as tape:
preds = model(x, training=True)
loss = loss_fn(y, preds)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
# first call traces the graph (slow), subsequent calls reuse it (fast)
train_step(x_batch, y_batch)
# experimental_relax_shapes allows retracing only when needed
@tf.function(reduce_retracing=True)
def fn(x): return model(x)
# input_signature pins the signature so no retracing
@tf.function(input_signature=[tf.TensorSpec([None, 784], tf.float32)])
def predict(x): return model(x)XLA Compilation
XLA (Accelerated Linear Algebra) compiles subgraphs into fused kernels, reducing launch overhead and enabling cross-op optimizations. jit_compile=True on a tf.function triggers XLA for that function. It's most effective for compute-heavy models (transformers, attention) and less helpful for I/O-bound ones. XLA requires static shapes within a compiled region — variable shapes cause recompilation. If you hit shape issues, pad to fixed shapes. Always benchmark: XLA can occasionally be slower for small models.
import tensorflow as tf
# JIT-compile a function with XLA (fuses ops for speed)
@tf.function(jit_compile=True)
def train_step(x, y):
with tf.GradientTape() as tape:
loss = loss_fn(y, model(x, training=True))
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
# enable XLA globally
tf.config.optimizer.set_jit(True)
# XLA fuses element-wise ops, reduces kernel launches,
# and can give 1.5-3x speedup on GPU/TPU
# especially effective for transformer / attention heavy modelsInput Pipeline Optimization
The bottleneck for GPU training is usually the input pipeline, not the model. The canonical pipeline caches after loading, shuffles, maps augmentation in parallel, batches, and prefetches. interleave reads multiple TFRecord files in parallel for disk-bound workloads. drop_remainder=True avoids a variable-shape last batch that forces graph retracing. Use the TF Profiler's Input Pipeline Analyzer to find your specific bottleneck. AUTOTUNE lets TF pick parallelism and prefetch sizes automatically.
import tensorflow as tf
# the canonical high-performance pipeline
ds = (tf.data.Dataset.from_tensor_slices((x, y))
.cache() # cache after loading
.shuffle(10000)
.map(augment, num_parallel_calls=tf.data.AUTOTUNE)
.batch(32, drop_remainder=True)
.prefetch(tf.data.AUTOTUNE))
# parallel file reading for TFRecord
files = tf.data.Dataset.list_files('data-*.tfrecord')
ds = files.interleave(
tf.data.TFRecordDataset,
cycle_length=4, num_parallel_calls=tf.data.AUTOTUNE)
# options for autotuning
options = tf.data.Options()
options.autotune.enabled = True
ds = ds.with_options(options)Mixed Precision & TF32
Mixed precision (float16/bfloat16) gives 2-3x throughput on Tensor Core GPUs. TF32 is even easier — it's a 19-bit mantissa format used transparently for matmul on Ampere+ (A100, H100), giving ~3x speedup with minimal accuracy loss and NO code changes or loss scaling. Enable both for maximum speed. TF32 is on by default in CUDA 11+ but TF exposes the flag for explicit control. For strict fp32 reproducibility (scientific computing), disable TF32.
import tensorflow as tf
# mixed precision (float16 compute, float32 master)
tf.keras.mixed_precision.set_global_policy('mixed_float16')
# enable TF32 on Ampere+ (big matmul speedup, slight precision loss)
tf.config.experimental.enable_tensor_float_32_execution(True)
# TF32 uses 19-bit mantissa for matmul on A100/H100
# it's transparent — no code changes, no loss scaling
# gives ~3x matmul speedup with minimal accuracy impact
# check policy
print(tf.keras.mixed_precision.global_policy())
print('TF32:', tf.config.experimental.tensor_float_32_execution_enabled())Graph Optimization
Grappler is TF's graph optimizer that runs automatically before execution: constant folding, layout optimization (NCHW vs NHWC for cuDNN), op fusion (remapping), dead code elimination, and more. These are all on by default — you only need to touch them to disable specific optimizations for debugging. debug_stripper removes Assert/CheckNumerics for production speed. Layout optimizer automatically picks NCHW for cuDNN convs on GPU even if your model specifies NHWC.
import tensorflow as tf
# enable grappler (default graph optimizer)
tf.config.optimizer.set_jit(True) # XLA
tf.config.optimizer.set_experimental_options({
'disable_model_pruning': False,
'disable_meta_optimizer': False,
'layout_optimizer': True, # optimize data layout (NCHW vs NHWC)
'constant_folding': True, # fold constants at build time
'shape_optimization': True,
'remapping': True, # fuse compatible ops
'arithmetic_optimization': True,
'loop_optimization': True,
'dependency_optimization': True,
'function_optimization': True,
'debug_stripper': True, # remove Assert/CheckNumerics
})
# all on by default — only set explicitly to disable specific onesAutoGraph & Control Flow
AutoGraph converts Python if/for/while inside tf.function into TF graph ops (tf.cond / tf.while_loop) so they run efficiently on GPU. Loops over a tf.data.Dataset with Python 'for' are fine — they iterate at trace time. Avoid if/for that depends on tensor VALUES (not shapes) — that forces retracing each time the value changes. For runtime control flow on tensor values, use tf.cond and tf.while_loop explicitly or let AutoGraph handle it (it will, but the conversion can be subtle).
import tensorflow as tf
# AutoGraph converts Python if/for/while into tf.cond / tf.while_loop
@tf.function
def fn(x):
if x > 0: # AutoGraph converts this
return x * 2
else:
return x
@tf.function
def sum_loop(n):
total = tf.constant(0)
for i in tf.range(n): # AutoGraph converts to tf.while_loop
total += i
return total
# Python-side (compile-time) loop vs TF-side (runtime) loop
@tf.function
def train_epoch(dataset):
for x, y in dataset: # Python iteration over dataset (fine)
train_step(x, y)
# avoid data-dependent Python control flow
# if x.shape[0] > 10: ... # OK (static shape)
# if tf.reduce_sum(x) > 0: ... # forces retracingСвязанные сниппеты TensorFlow
Copy-paste ready code for common tasks.
Tensor Basics
Create and operate on TensorFlow tensors.
Keras Model
Build models with Sequential and the functional API.
Layers
Use core layers and build a custom one.
Compile and Train
Compile, fit, and evaluate a Keras model.
Custom Training Loop
Step through batches with GradientTape.
Callbacks
Monitor and control training with callbacks.
Save and Load
Persist models in SavedModel and Keras formats.
Data Pipeline
Build efficient input pipelines with tf.data.
Was this helpful?