15. Custom Layers, Models and Training Loops¶
Move beyond the standard Keras
Sequentialand Functional APIs and learn how to build custom layers, custom models, and custom training loops with TensorFlow and Keras. Understand how forward computation, trainable parameters, automatic differentiation, gradient computation, and optimizer updates work together to provide fine-grained control over Deep Learning systems.
๐ฏ Learning Objectives¶
After completing this chapter, you will be able to:
- Understand why custom Keras components are required
- Create custom Keras layers
- Understand the
Layerbase class - Implement the
build()method - Implement the
call()method - Create trainable and non-trainable variables
- Understand
trainable_weights - Understand
non_trainable_weights - Build custom mathematical operations as layers
- Create reusable custom layers
- Create custom Keras models
- Understand the difference between custom layers and custom models
- Implement custom forward passes
- Understand
training=Trueandtraining=False - Use
tf.GradientTape - Compute gradients manually
- Apply gradients using an optimizer
- Implement a custom training step
- Build a custom training loop
- Implement validation loops
- Track metrics manually
- Use
tf.function - Understand eager execution vs graph execution
- Understand how Keras
fit()works conceptually - Override
train_step() - Combine custom training logic with
model.fit() - Implement custom losses
- Implement custom metrics
- Implement custom regularization
- Understand when custom training loops are appropriate
- Design maintainable custom Deep Learning components
๐ Overview¶
Keras provides high-level APIs that make Deep Learning development straightforward:
For many applications, these abstractions are sufficient.
However, advanced Deep Learning systems sometimes require more control.
Examples include:
Custom Layers
Custom Forward Pass
Custom Loss
Custom Optimization
Multiple Gradient Computations
Gradient Manipulation
Custom Training Logic
Adversarial Training
Contrastive Learning
Multi-Task Optimization
Research Architectures
For these situations, Keras allows developers to move down the abstraction stack.
flowchart TD
HIGH["High-Level Keras API"]
HIGH --> SEQ["Sequential"]
HIGH --> FUNC["Functional API"]
HIGH --> FIT["model.fit()"]
CUSTOM["Custom Keras Components"]
CUSTOM --> LAYER["Custom Layers"]
CUSTOM --> MODEL["Custom Models"]
CUSTOM --> TRAIN["Custom Training Loops"]
LAYER --> LOW["Lower-Level TensorFlow Control"]
MODEL --> LOW
TRAIN --> LOW
LOW --> GRAD["GradientTape"]
LOW --> OPT["Optimizer"]
LOW --> TENSOR["Tensor Operations"]๐ง Why Custom Layers?¶
Standard Keras already provides many layers:
But sometimes your architecture requires an operation that does not exist as a standard layer.
For example:
Custom Mathematical Transformation
Custom Feature Interaction
Custom Normalization
Custom Attention
Custom Routing
Custom Residual Block
Custom Research Layer
Instead of writing the operation directly inside the model, create a reusable layer.
๐งฑ Custom Layer Concept¶
A custom layer can be viewed as:
flowchart LR
INPUT["Input Tensor"]
LAYER["Custom Layer"]
TRANSFORM["Custom Computation"]
OUTPUT["Output Tensor"]
INPUT --> LAYER
LAYER --> TRANSFORM
TRANSFORM --> OUTPUT๐ง tf.keras.layers.Layer¶
Custom layers generally inherit from:
Basic structure:
class MyLayer(tf.keras.layers.Layer):
def __init__(self, ...):
super().__init__()
def build(self, input_shape):
...
def call(self, inputs):
...
The three important concepts are:
๐ง Role of __init__()¶
The constructor is generally used to configure the layer.
For example:
class MyLayer(tf.keras.layers.Layer):
def __init__(
self,
units
):
super().__init__()
self.units = units
The constructor stores configuration.
It should not necessarily create weights that depend on the input shape.
๐ง Role of build()¶
build() is useful when weights depend on the input shape.
Example:
def build(
self,
input_shape
):
self.w = self.add_weight(
shape=(
input_shape[-1],
self.units
),
initializer="random_normal",
trainable=True
)
self.b = self.add_weight(
shape=(self.units,),
initializer="zeros",
trainable=True
)
This allows the layer to determine the required weight dimensions after the input shape is known.
๐ง Role of call()¶
call() defines the forward computation.
For example:
Conceptually:
๐งช Complete Custom Dense Layer¶
import tensorflow as tf
class CustomDense(
tf.keras.layers.Layer
):
def __init__(
self,
units
):
super().__init__()
self.units = units
def build(
self,
input_shape
):
self.w = self.add_weight(
shape=(
input_shape[-1],
self.units
),
initializer="random_normal",
trainable=True
)
self.b = self.add_weight(
shape=(
self.units,
),
initializer="zeros",
trainable=True
)
def call(
self,
inputs
):
return tf.matmul(
inputs,
self.w
) + self.b
Usage:
๐งฎ Custom Dense Layer Mathematics¶
The custom layer implements:
[ y = Wx+b ]
If an activation is included:
[ y=f(Wx+b) ]
This demonstrates an important principle:
A neural-network layer is fundamentally a parameterized mathematical transformation.
๐ง add_weight()¶
Keras provides:
for creating layer variables.
Example:
self.kernel = self.add_weight(
shape=(input_dim, units),
initializer="glorot_uniform",
trainable=True
)
This allows Keras to track the variable automatically.
๐ Trainable Variables¶
A layer exposes trainable parameters through:
For a custom dense layer:
๐ง Trainable vs Non-Trainable Variables¶
A layer may contain:
Trainable variables are updated during optimization.
Non-trainable variables are not updated by gradient descent.
flowchart TD
LAYER["Keras Layer"]
LAYER --> TRAIN["Trainable Variables"]
LAYER --> NON["Non-Trainable Variables"]
TRAIN --> GRAD["Gradients"]
GRAD --> UPDATE["Optimizer Update"]
NON --> STATE["Layer State"]๐งช Non-Trainable Weight¶
This variable is tracked by Keras but is not optimized through backpropagation.
๐ง Custom Layer with Activation¶
class CustomDense(
tf.keras.layers.Layer
):
def __init__(
self,
units,
activation=None
):
super().__init__()
self.units = units
self.activation = tf.keras.activations.get(
activation
)
def build(
self,
input_shape
):
self.w = self.add_weight(
shape=(
input_shape[-1],
self.units
),
initializer="glorot_uniform"
)
self.b = self.add_weight(
shape=(self.units,),
initializer="zeros"
)
def call(
self,
inputs
):
output = tf.matmul(
inputs,
self.w
) + self.b
if self.activation is not None:
output = self.activation(
output
)
return output
๐ง Why build() Is Useful¶
Consider:
At this moment, the layer may not know:
When the first input arrives:
This makes custom layers reusable across different input dimensions.
๐ Layer Lifecycle¶
flowchart TD
CREATE["Create Layer"]
INIT["__init__()"]
INPUT["First Input"]
BUILD["build()"]
CALL["call()"]
OUTPUT["Output"]
CREATE --> INIT
INIT --> INPUT
INPUT --> BUILD
BUILD --> CALL
CALL --> OUTPUTFor subsequent compatible calls, Keras generally reuses the already-created weights rather than rebuilding them.
๐ง Custom Layer Example โ Scaling¶
A simple custom layer can implement:
[ y=\alpha x ]
class ScalingLayer(
tf.keras.layers.Layer
):
def __init__(
self,
scale
):
super().__init__()
self.scale = scale
def call(
self,
inputs
):
return inputs * self.scale
Usage:
๐ง Custom Layer Example โ Learnable Scaling¶
The scale itself can be trainable.
class LearnableScaling(
tf.keras.layers.Layer
):
def build(
self,
input_shape
):
self.scale = self.add_weight(
shape=(input_shape[-1],),
initializer="ones",
trainable=True
)
def call(
self,
inputs
):
return inputs * self.scale
The model learns:
๐ง Custom Layer Example โ Residual Block¶
A residual block can be implemented as a custom layer.
class ResidualBlock(
tf.keras.layers.Layer
):
def __init__(
self,
units
):
super().__init__()
self.dense1 = tf.keras.layers.Dense(
units,
activation="relu"
)
self.dense2 = tf.keras.layers.Dense(
units
)
self.activation = tf.keras.layers.ReLU()
def call(
self,
inputs
):
x = self.dense1(
inputs
)
x = self.dense2(
x
)
x = x + inputs
return self.activation(
x
)
Architecture:
flowchart TD
INPUT["Input"]
D1["Dense + ReLU"]
D2["Dense"]
ADD["Add"]
RELU["ReLU"]
INPUT --> D1
D1 --> D2
D2 --> ADD
INPUT --> ADD
ADD --> RELUThis pattern becomes important for:
- ResNet
- Deep MLPs
- Transformer blocks
- Advanced architectures
๐ง Custom Layer with Training Behavior¶
Some layers behave differently during training and inference.
Examples:
The call() method can accept:
Example:
class CustomDropout(
tf.keras.layers.Layer
):
def __init__(
self,
rate
):
super().__init__()
self.rate = rate
def call(
self,
inputs,
training=False
):
if training:
return tf.nn.dropout(
inputs,
rate=self.rate
)
return inputs
๐ง Training vs Inference¶
flowchart TD
INPUT["Input"]
TRAIN{"training=True?"}
TRAIN -->|Yes| TRAINPATH["Training Behavior"]
TRAIN -->|No| INFER["Inference Behavior"]
TRAINPATH --> OUTPUT["Output"]
INFER --> OUTPUTThe model should be explicit about behavior that differs between training and inference.
๐ง Custom Models¶
A custom layer represents a reusable transformation.
A custom model generally represents the larger architecture.
For example:
๐งช Custom Keras Model¶
class MyModel(
tf.keras.Model
):
def __init__(
self
):
super().__init__()
self.dense1 = tf.keras.layers.Dense(
128,
activation="relu"
)
self.dense2 = tf.keras.layers.Dense(
64,
activation="relu"
)
self.output_layer = tf.keras.layers.Dense(
10,
activation="softmax"
)
def call(
self,
inputs,
training=False
):
x = self.dense1(
inputs
)
x = self.dense2(
x
)
return self.output_layer(
x
)
Usage:
๐ง Layer vs Model¶
| Custom Layer | Custom Model |
|---|---|
| Reusable computation block | Complete architecture |
| Usually smaller | Usually larger |
Extends Layer |
Extends Model |
| Can contain other layers | Can contain many layers |
| Used inside models | Represents the model itself |
| Defines transformation | Defines overall forward computation |
๐ Custom Model Architecture¶
flowchart LR
INPUT["Input"]
L1["Custom / Standard Layer"]
L2["Custom / Standard Layer"]
HEAD["Output Head"]
OUTPUT["Prediction"]
INPUT --> L1
L1 --> L2
L2 --> HEAD
HEAD --> OUTPUT๐ง call() Defines Forward Pass¶
For a custom model:
defines how data moves through the network.
Conceptually:
[ \hat{y}=f_\theta(x) ]
where:
๐ง Automatic Differentiation¶
Training requires gradients.
TensorFlow provides:
for automatic differentiation.
Conceptually:
flowchart TD
INPUT["Input"]
MODEL["Model"]
PRED["Prediction"]
LOSS["Loss"]
TAPE["GradientTape"]
GRAD["Gradients"]
OPT["Optimizer"]
UPDATE["Updated Parameters"]
INPUT --> MODEL
MODEL --> PRED
PRED --> LOSS
LOSS --> TAPE
TAPE --> GRAD
GRAD --> OPT
OPT --> UPDATE
UPDATE --> MODEL๐งฎ Gradient Computation¶
Suppose:
[ L=f(w) ]
The gradient is:
[ \frac{\partial L}{\partial w} ]
The optimizer then updates:
[ w \leftarrow w-\eta \frac{\partial L}{\partial w} ]
๐งช Basic GradientTape¶
x = tf.Variable(
3.0
)
with tf.GradientTape() as tape:
y = x ** 2
gradient = tape.gradient(
y,
x
)
print(
gradient
)
Because:
[ y=x^2 ]
we have:
[ \frac{dy}{dx}=2x ]
For:
[ x=3 ]
the gradient is:
๐ง GradientTape Workflow¶
flowchart LR
VARIABLE["Trainable Variable"]
TAPE["GradientTape"]
FORWARD["Forward Computation"]
LOSS["Loss"]
GRAD["Gradient"]
VARIABLE --> TAPE
TAPE --> FORWARD
FORWARD --> LOSS
LOSS --> GRADThe tape records operations performed within its context so gradients can later be computed.
๐งช Simple Linear Regression with GradientTape¶
import tensorflow as tf
w = tf.Variable(
0.0
)
b = tf.Variable(
0.0
)
x = tf.constant(
[1.0, 2.0, 3.0]
)
y_true = tf.constant(
[2.0, 4.0, 6.0]
)
with tf.GradientTape() as tape:
y_pred = w * x + b
loss = tf.reduce_mean(
tf.square(
y_true - y_pred
)
)
gradients = tape.gradient(
loss,
[w, b]
)
print(
gradients
)
This demonstrates the core mechanism behind neural-network training.
๐ง Applying Gradients¶
Once gradients are computed:
The complete process becomes:
๐งช Manual Training Step¶
optimizer = tf.keras.optimizers.Adam(
learning_rate=0.001
)
with tf.GradientTape() as tape:
predictions = model(
x_batch,
training=True
)
loss = loss_fn(
y_batch,
predictions
)
gradients = tape.gradient(
loss,
model.trainable_variables
)
optimizer.apply_gradients(
zip(
gradients,
model.trainable_variables
)
)
This is the fundamental custom training step.
๐ง Custom Training Loop¶
A custom training loop repeatedly executes the training step.
for epoch in range(
epochs
):
for x_batch, y_batch in train_dataset:
with tf.GradientTape() as tape:
predictions = model(
x_batch,
training=True
)
loss = loss_fn(
y_batch,
predictions
)
gradients = tape.gradient(
loss,
model.trainable_variables
)
optimizer.apply_gradients(
zip(
gradients,
model.trainable_variables
)
)
๐ Custom Training Loop¶
flowchart TD
DATA["Training Dataset"]
EPOCH["Epoch"]
BATCH["Batch"]
FORWARD["Forward Pass"]
LOSS["Loss"]
GRAD["GradientTape"]
UPDATE["Optimizer Update"]
NEXT["Next Batch"]
DATA --> EPOCH
EPOCH --> BATCH
BATCH --> FORWARD
FORWARD --> LOSS
LOSS --> GRAD
GRAD --> UPDATE
UPDATE --> NEXT
NEXT --> BATCH๐ง Why Use Custom Training Loops?¶
Use custom training loops when you need control over:
- Multiple losses
- Multiple optimizers
- Custom gradient manipulation
- Gradient accumulation
- Adversarial training
- Contrastive learning
- Custom optimization logic
- Complex training schedules
- Specialized research architectures
For ordinary supervised learning, model.fit() is often simpler.
๐ง model.fit() vs Custom Loop¶
model.fit() |
Custom Training Loop |
|---|---|
| High-level | Low-level |
| Less code | More code |
| Standard workflows | Specialized workflows |
| Built-in callbacks | Manual control |
| Built-in metrics | Manual or hybrid |
| Easier maintenance | More responsibility |
| Faster development | Greater flexibility |
๐ง Keras fit() Internals¶
Conceptually, model.fit() performs something similar to:
for each epoch:
for each batch:
forward pass
calculate loss
calculate gradients
update parameters
update metrics
The difference is that Keras handles the training infrastructure for you.
๐ง Overriding train_step()¶
Keras provides a useful middle ground.
Instead of implementing the entire training loop, override:
This allows custom training logic while still using:
๐งช Custom train_step()¶
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.compute_loss(
x=x,
y=y,
y_pred=y_pred
)
gradients = tape.gradient(
loss,
self.trainable_variables
)
self.optimizer.apply_gradients(
zip(
gradients,
self.trainable_variables
)
)
return {
"loss": loss
}
This provides custom training behavior without completely abandoning the Keras training framework.
๐ง Three Levels of Training Control¶
flowchart TD
HIGH["High Control Spectrum"]
HIGH --> FIT["model.fit()"]
FIT --> TRAINSTEP["Override train_step()"]
TRAINSTEP --> LOOP["Fully Custom Training Loop"]
LOOP --> LOW["Maximum Control"]Think of the choices as:
Standard Training
โ
model.fit()
Custom Training Behavior
โ
Custom train_step()
Maximum Control
โ
Custom Training Loop
๐ง Custom Loss Functions¶
Keras supports custom losses.
A loss function generally has:
Example:
๐งฎ Mean Squared Error¶
[ MSE = \frac{1}{n} \sum_{i=1}^{n} (y_i-\hat{y}_i)^2 ]
Custom losses are useful when standard loss functions do not match the training objective.
๐ง Custom Loss with Additional Terms¶
A custom loss can combine multiple objectives.
For example:
[ L = L_{task} + \lambda L_{regularization} ]
This is common in:
- Representation learning
- Multi-task learning
- Regularized models
- Research architectures
๐ง Custom Regularization¶
A layer can add a regularization term using:
Example:
class RegularizedLayer(
tf.keras.layers.Layer
):
def call(
self,
inputs
):
penalty = tf.reduce_mean(
tf.square(inputs)
)
self.add_loss(
1e-4 * penalty
)
return inputs
Keras then incorporates this additional loss during training.
๐ง add_loss()¶
Conceptually:
flowchart TD
TASK["Task Loss"]
REG["Regularization Loss"]
TASK --> TOTAL["Total Loss"]
REG --> TOTAL
TOTAL --> GRAD["Gradient"]
GRAD --> UPDATE["Optimizer"]๐ง Custom Metrics¶
Metrics measure model behavior without necessarily controlling optimization.
Example:
class MeanAbsoluteErrorMetric(
tf.keras.metrics.Metric
):
def __init__(
self,
name="mae",
**kwargs
):
super().__init__(
name=name,
**kwargs
)
self.total = self.add_weight(
name="total",
initializer="zeros"
)
self.count = self.add_weight(
name="count",
initializer="zeros"
)
def update_state(
self,
y_true,
y_pred,
sample_weight=None
):
error = tf.abs(
y_true - y_pred
)
self.total.assign_add(
tf.reduce_sum(error)
)
self.count.assign_add(
tf.cast(
tf.size(error),
tf.float32
)
)
def result(self):
return (
self.total /
self.count
)
def reset_state(self):
self.total.assign(
0.0
)
self.count.assign(
0.0
)
๐ง Loss vs Metric¶
Keep the distinction clear:
For example:
The optimizer uses gradients derived from the loss.
๐ง Multiple Optimizers¶
Custom training loops can use different optimizers for different components.
For example:
This pattern is important in:
- GANs
- Adversarial training
- Multi-network architectures
๐ง Multiple Gradient Tapes¶
Some architectures require separate gradient computations.
Conceptually:
flowchart TD
INPUT["Input"]
MODEL_A["Model A"]
MODEL_B["Model B"]
LOSS_A["Loss A"]
LOSS_B["Loss B"]
GRAD_A["Gradients A"]
GRAD_B["Gradients B"]
OPT_A["Optimizer A"]
OPT_B["Optimizer B"]
INPUT --> MODEL_A
INPUT --> MODEL_B
MODEL_A --> LOSS_A
MODEL_B --> LOSS_B
LOSS_A --> GRAD_A
LOSS_B --> GRAD_B
GRAD_A --> OPT_A
GRAD_B --> OPT_BThis level of control is difficult to express using a simple standard training workflow.
๐ง Gradient Clipping in Custom Loops¶
Gradients can be clipped before applying them.
gradients = tape.gradient(
loss,
model.trainable_variables
)
gradients, _ = tf.clip_by_global_norm(
gradients,
1.0
)
optimizer.apply_gradients(
zip(
gradients,
model.trainable_variables
)
)
The workflow becomes:
๐ง Gradient Accumulation¶
Gradient accumulation can simulate a larger effective batch size when GPU memory is limited.
Conceptually:
Batch 1 โ Gradients
Batch 2 โ Accumulate
Batch 3 โ Accumulate
Batch 4 โ Accumulate
โ
Update Weights
flowchart LR
B1["Batch 1"]
B2["Batch 2"]
B3["Batch 3"]
B4["Batch 4"]
B1 --> G["Gradient Accumulator"]
B2 --> G
B3 --> G
B4 --> G
G --> UPDATE["Optimizer Update"]This is useful for large models where the desired effective batch size does not fit into GPU memory.
๐งช Simplified Gradient Accumulation¶
accumulated_gradients = [
tf.zeros_like(variable)
for variable
in model.trainable_variables
]
for x_batch, y_batch in train_dataset:
with tf.GradientTape() as tape:
predictions = model(
x_batch,
training=True
)
loss = loss_fn(
y_batch,
predictions
)
gradients = tape.gradient(
loss,
model.trainable_variables
)
accumulated_gradients = [
acc + grad
for acc, grad
in zip(
accumulated_gradients,
gradients
)
]
In production, the accumulated gradients should be normalized appropriately and reset after the optimizer update.
๐ง tf.function¶
TensorFlow can convert Python functions into optimized TensorFlow graphs using:
Example:
This can improve execution efficiency for repeated TensorFlow operations.
๐ง Eager Execution vs Graph Execution¶
Eager Execution¶
Operations execute immediately.
Graph Execution¶
Operations are traced into a computation graph.
flowchart LR
EAGER["Eager Execution"]
GRAPH["Graph Execution"]
EAGER --> DEBUG["Easy Debugging"]
GRAPH --> OPT["Graph Optimization"]๐ง tf.function Example¶
@tf.function
def train_step(
model,
optimizer,
loss_fn,
x,
y
):
with tf.GradientTape() as tape:
predictions = model(
x,
training=True
)
loss = loss_fn(
y,
predictions
)
gradients = tape.gradient(
loss,
model.trainable_variables
)
optimizer.apply_gradients(
zip(
gradients,
model.trainable_variables
)
)
return loss
This can be used inside a custom training loop.
โ tf.function Considerations¶
tf.function can improve performance, but developers should understand:
- Tracing
- Retracing
- Python side effects
- Tensor vs Python values
- Dynamic control flow
- Input signatures
Avoid repeatedly creating different Python signatures that cause unnecessary retracing.
๐ง Custom Validation Loop¶
A custom validation loop can be implemented without gradient computation.
for x_batch, y_batch in validation_dataset:
predictions = model(
x_batch,
training=False
)
loss = loss_fn(
y_batch,
predictions
)
Notice:
and:
because validation does not update model parameters.
๐ง Training vs Validation¶
flowchart TD
BATCH["Batch"]
BATCH --> TRAIN{"Training?"}
TRAIN -->|Yes| FORWARD["Forward Pass"]
FORWARD --> LOSS["Loss"]
LOSS --> TAPE["GradientTape"]
TAPE --> GRAD["Gradients"]
GRAD --> UPDATE["Optimizer Update"]
TRAIN -->|No| EVAL["Forward Pass Only"]
EVAL --> METRIC["Metrics"]๐ง Custom Training Step with Metrics¶
A more complete training step can track metrics.
loss_metric = tf.keras.metrics.Mean(
name="loss"
)
for x_batch, y_batch in train_dataset:
with tf.GradientTape() as tape:
predictions = model(
x_batch,
training=True
)
loss = loss_fn(
y_batch,
predictions
)
gradients = tape.gradient(
loss,
model.trainable_variables
)
optimizer.apply_gradients(
zip(
gradients,
model.trainable_variables
)
)
loss_metric.update_state(
loss
)
At the end:
๐ง Complete Custom Training Architecture¶
flowchart TD
DATA["Training Dataset"]
EPOCH["Epoch"]
BATCH["Batch"]
MODEL["Custom Model"]
LOSS["Custom / Standard Loss"]
TAPE["GradientTape"]
GRAD["Gradients"]
CLIP["Gradient Processing"]
OPT["Optimizer"]
METRIC["Metrics"]
UPDATE["Updated Parameters"]
DATA --> EPOCH
EPOCH --> BATCH
BATCH --> MODEL
MODEL --> LOSS
LOSS --> TAPE
TAPE --> GRAD
GRAD --> CLIP
CLIP --> OPT
OPT --> UPDATE
UPDATE --> MODEL
LOSS --> METRIC
MODEL --> METRIC๐ง Custom Training Loop Design Principles¶
A good custom training loop should have clear separation between:
Data Loading
โ
Forward Pass
โ
Loss Calculation
โ
Gradient Calculation
โ
Gradient Processing
โ
Parameter Update
โ
Metrics
โ
Checkpointing
Avoid putting all of this logic into one giant function.
๐งฑ Recommended Structure¶
A production-oriented implementation could look like:
training/
โ
โโโ trainer.py
โโโ losses.py
โโโ metrics.py
โโโ optimizers.py
โโโ schedulers.py
โโโ callbacks.py
models/
โ
โโโ layers.py
โโโ blocks.py
โโโ model.py
data/
โ
โโโ dataset.py
โโโ preprocessing.py
This keeps training infrastructure separate from model architecture.
๐ง Custom Layer Testing¶
Custom layers should be tested independently.
For example:
layer = CustomDense(
32
)
x = tf.random.normal(
(8, 16)
)
y = layer(
x
)
assert y.shape == (
8,
32
)
Test:
Output Shape
Data Type
Trainable Variables
Numerical Behavior
Serialization
Training Behavior
Inference Behavior
๐งช Numerical Gradient Verification¶
For complex custom operations, gradient correctness matters.
You can inspect gradients:
with tf.GradientTape() as tape:
output = layer(x)
gradient = tape.gradient(
output,
layer.trainable_variables
)
Check for:
โ Common Custom-Layer Errors¶
Common mistakes include:
- Creating weights in the wrong place
- Creating weights repeatedly inside
call() - Forgetting
trainable=True - Returning tensors with unexpected shapes
- Ignoring the
trainingargument - Mixing NumPy operations with TensorFlow tensors inside differentiable computation
- Creating Python-side state that TensorFlow cannot track
- Failing to implement serialization configuration
- Assuming eager execution behavior is identical under
tf.function
๐ง Serialization of Custom Layers¶
Custom layers should support serialization when they need to be saved and reloaded.
Example:
class ScalingLayer(
tf.keras.layers.Layer
):
def __init__(
self,
scale,
**kwargs
):
super().__init__(
**kwargs
)
self.scale = scale
def call(
self,
inputs
):
return inputs * self.scale
def get_config(
self
):
config = super().get_config()
config.update({
"scale": self.scale
})
return config
This allows Keras to reconstruct the layer configuration.
๐ง Why get_config() Matters¶
Without proper serialization:
With serialization support:
For production systems, model portability is an important consideration.
๐ง Custom Models and Serialization¶
Custom models should similarly be designed with serialization in mind.
Keep configuration values explicit:
class CustomModel(
tf.keras.Model
):
def __init__(
self,
units,
**kwargs
):
super().__init__(
**kwargs
)
self.units = units
self.dense = tf.keras.layers.Dense(
units
)
Configuration should not be hidden inside arbitrary runtime state.
๐ข Enterprise Perspective¶
Custom layers and training loops provide powerful capabilities, but they also increase engineering responsibility.
With standard Keras:
With custom training:
The production engineering team must now consider:
Correctness
+
Testing
+
Serialization
+
Reproducibility
+
Performance
+
Numerical Stability
+
Monitoring
+
Distributed Execution
Production Insight
Use the highest-level abstraction that satisfies the requirement.
A practical decision is:
Standard Model
โ
model.fit()
Need Custom Model Component
โ
Custom Layer
Need Custom Training Behavior
โ
Override train_step()
Need Maximum Training Control
โ
Custom Training Loop
Do not implement a custom training loop simply because you can.
Every additional layer of customization becomes part of the production system that must be tested, maintained, monitored, and upgraded.
๐ง Custom Training Decision Guide¶
flowchart TD
START["Training Requirement"]
START --> STANDARD{"Standard Supervised Training?"}
STANDARD -->|Yes| FIT["Use model.fit()"]
STANDARD -->|No| COMPONENT{"Need Custom Layer / Operation?"}
COMPONENT -->|Yes| LAYER["Create Custom Layer"]
COMPONENT -->|No| TRAINSTEP{"Need Custom Training Logic?"}
TRAINSTEP -->|Moderate| OVERRIDE["Override train_step()"]
TRAINSTEP -->|Extensive| LOOP["Custom Training Loop"]๐งช Practical Exercise 1 โ Custom Dense Layer¶
Implement a layer that performs:
[ y=f(Wx+b) ]
Requirements:
Test it against:
for compatible initialization and configuration.
๐งช Practical Exercise 2 โ Custom Residual Block¶
Implement:
Input
โ
โโโโโโโโโโโโโโโโ
โ โ
Dense โ
โ โ
Dense โ
โ โ
โโโโโ Add โโโโโโ
โ
ReLU
Verify:
๐งช Practical Exercise 3 โ Custom Loss¶
Implement:
Compare it with:
Verify that the outputs are numerically consistent for the same inputs.
๐งช Practical Exercise 4 โ Custom Training Loop¶
Build:
Track:
๐งช Practical Exercise 5 โ Custom train_step()¶
Create a custom model that overrides:
but still trains using:
Compare the implementation with a normal Keras model.
๐งช Practical Exercise 6 โ Gradient Clipping¶
Modify the custom training loop to apply:
Compare:
Track gradient norms and training stability.
๐งช Practical Exercise 7 โ Gradient Accumulation¶
Implement gradient accumulation:
Compare it against:
with approximately the same effective batch size.
๐งช Practical Exercise 8 โ Training vs Inference¶
Create a custom layer whose behavior changes according to:
and:
Verify that:
and:
behave as expected.
๐ง Interview Questions¶
Beginner¶
1. Why would you create a custom Keras layer?¶
To implement reusable computations or transformations that are not adequately represented by existing Keras layers.
2. What are the main methods in a custom layer?¶
Commonly:
3. What does build() do?¶
It is commonly used to create weights whose shapes depend on the input shape.
4. What does call() do?¶
It defines the forward computation performed by the layer.
5. What is GradientTape?¶
tf.GradientTape records differentiable TensorFlow operations so gradients can be computed automatically.
Intermediate¶
6. What is the difference between a custom layer and custom model?¶
A custom layer generally represents a reusable computation block, while a custom model represents a larger network architecture and its forward pass.
7. What is a custom training loop?¶
It is a manually controlled training process where the developer explicitly performs forward computation, loss calculation, gradient calculation, and parameter updates.
8. Why use a custom training loop?¶
When standard model.fit() does not provide sufficient control over the training algorithm.
9. What is train_step()?¶
It is a Keras extension point that allows custom training logic while continuing to use the broader model.fit() framework.
10. What is the difference between a loss and a metric?¶
A loss provides the optimization objective, while a metric is primarily used to measure and monitor model behavior.
11. What is gradient clipping?¶
Gradient clipping limits gradient magnitude to improve numerical stability and control exploding gradients.
Advanced¶
12. Why should weights generally not be created inside call()?¶
Because call() may execute repeatedly. Creating weights there can lead to repeated variable creation and incorrect parameter tracking.
13. Why is training passed to call()?¶
Some layers need different behavior during training and inference, such as Dropout and Batch Normalization.
14. What is the advantage of overriding train_step() over writing a completely custom loop?¶
It allows custom training logic while retaining Keras features such as fit(), callbacks, progress reporting, and other training infrastructure.
15. When would you use multiple optimizers?¶
When different model components require independent optimization, such as generator/discriminator systems or specialized multi-network training.
16. Why is gradient accumulation useful?¶
It allows a larger effective batch size when the desired batch cannot fit into available device memory.
17. What are the risks of custom training loops?¶
They introduce additional complexity around:
Correctness
Metrics
Checkpointing
Distributed Training
Mixed Precision
Serialization
Reproducibility
18. Why is serialization important for custom layers?¶
A production model must be reconstructable after deployment or loading. Custom layer configuration must therefore be preserved correctly.
19. When should you avoid a custom training loop?¶
When standard Keras training already satisfies the requirements. Unnecessary customization increases maintenance and testing complexity.
20. How would you design a production custom training system?¶
Separate:
and ensure each component is testable and observable.
๐ Key Takeaways¶
- Keras provides high-level APIs but also allows low-level customization.
- Custom layers extend
tf.keras.layers.Layer. __init__()generally stores configuration.build()is useful for creating input-dependent weights.call()defines the forward computation.add_weight()allows Keras to track variables correctly.- Trainable variables participate in gradient-based optimization.
- Non-trainable variables represent state that is not optimized through gradients.
- Custom layers can implement reusable mathematical operations.
- Custom models extend
tf.keras.Model. - A model's
call()defines its forward pass. tf.GradientTapeprovides automatic differentiation.- A custom training step generally consists of forward pass, loss calculation, gradient computation, and optimizer update.
- Custom training loops provide maximum control.
- Overriding
train_step()provides a useful middle ground betweenmodel.fit()and a fully custom loop. - Custom losses allow specialized optimization objectives.
add_loss()can incorporate additional regularization or auxiliary objectives.- Custom metrics provide specialized monitoring.
- Gradient clipping can be incorporated into custom training.
- Gradient accumulation can increase effective batch size without increasing physical batch size.
tf.functioncan convert Python functions into optimized TensorFlow graphs.- Training and inference may require different layer behavior.
- Custom components should be independently tested.
- Custom layers and models should support serialization when required.
- Custom training loops should be used only when the additional control justifies the additional complexity.
- Production Deep Learning requires balancing flexibility with maintainability, reproducibility, observability, and operational reliability.
๐ Further Reading¶
Continue with:
- 16. PyTorch Fundamentals and Tensors
- 17. PyTorch Autograd, Dataset and DataLoader
- 18. Building Classification and Regression Models
- 19. Convolutional Neural Networks
- 20. CNN Architecture, Optimization and Training
- 22. ResNet, Residual Connections and TorchVision
- 30. Generative Adversarial Networks
- 36. Deep Learning Training and Model Lifecycle
The next chapter moves to the second major Deep Learning framework and builds the equivalent foundations in PyTorch, including tensors, autograd, model construction, and GPU execution.
โก๏ธ Next Chapter¶
16. PyTorch Fundamentals and Tensors
Enterprise AI Engineering Handbook
Building Production-Grade Enterprise AI Systems โ One Chapter at a Time.