TensorFlow
Keras
Machine Learning
Deep Learning
Neural Networks

TensorFlow - Difference between tf.keras.layers.Layer vs tf.keras.Model

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

tf.keras.Model is a subclass of tf.keras.layers.Layer, so they share a lot of behavior. The practical difference is intent: use Layer for reusable building blocks and use Model for an object that represents a trainable end-to-end network with training, evaluation, saving, and summary features.

What a Layer Is

A Layer encapsulates computation plus state such as weights. It is the right abstraction for things like:

  • a custom attention block
  • a residual block
  • a normalization or embedding component
  • any reusable transformation inside a larger network

Example:

python
1import tensorflow as tf
2
3
4class ScaleLayer(tf.keras.layers.Layer):
5    def build(self, input_shape):
6        self.scale = self.add_weight(
7            shape=(1,),
8            initializer="ones",
9            trainable=True,
10        )
11
12    def call(self, inputs):
13        return inputs * self.scale

This is a self-contained component, but it is not necessarily the whole model.

What a Model Adds

A Model is still a Layer, but it is treated as a top-level network. That means it provides conveniences such as:

  • 'fit'
  • 'evaluate'
  • 'predict'
  • 'save'
  • 'summary'

Here is a small subclassed model that uses the custom layer above:

python
1import tensorflow as tf
2
3
4class SimpleModel(tf.keras.Model):
5    def __init__(self):
6        super().__init__()
7        self.scale = ScaleLayer()
8        self.dense = tf.keras.layers.Dense(1)
9
10    def call(self, inputs):
11        x = self.scale(inputs)
12        return self.dense(x)
13
14
15model = SimpleModel()
16model.compile(optimizer="adam", loss="mse")

This object is now the trainable model, not just one internal transformation.

A Helpful Mental Model

Think of it this way:

  • every Model is a Layer
  • not every Layer should be a Model

If you subclass Model for every little component, you blur the boundary between "reusable part" and "trainable artifact." That makes code harder to understand and can create confusion around saving, summaries, and nesting. It also makes testing individual building blocks less clear because their role in the architecture is no longer obvious.

Can Models Be Nested

Yes. Since Model inherits from Layer, one model can be used inside another. That is useful for encoder-decoder systems, shared submodels, or pretrained backbones.

But even when nesting is valid, the abstraction still matters. If a block is just a block, subclassing Layer usually communicates intent better.

When to Choose Which

Choose Layer when:

  • the object is a component inside a bigger model
  • you want reuse across several architectures
  • the object does not need to act like the primary training artifact

Choose Model when:

  • the object is the network you want to compile and train
  • you want summary, fit, and saving behavior at that boundary
  • the object is the natural top-level artifact for inference or export

Functional API Note

If you build with the Functional API, Keras already creates a Model object for you:

python
inputs = tf.keras.Input(shape=(4,))
outputs = tf.keras.layers.Dense(1)(inputs)
model = tf.keras.Model(inputs, outputs)

The same principle still holds. Layers are parts. A model is the assembled network boundary.

Common Pitfalls

The biggest mistake is subclassing Model for every custom operation. That works technically, but it makes the codebase noisier than necessary.

Another mistake is forgetting that Layer can own weights too. Some developers incorrectly assume only Model can track trainable variables.

A third issue is choosing Layer for the real top-level network and then reimplementing training and saving behavior manually when Model already provides it.

Summary

  • 'tf.keras.Model inherits from tf.keras.layers.Layer.'
  • Use Layer for reusable building blocks with state and computation.
  • Use Model for the top-level trainable network boundary.
  • Models can be nested because they are also layers.
  • The best choice is about intent and API boundary, not raw capability alone.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.