CINXE.COM

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/guides/making_new_layers_and_models_via_subclassing/" />  <meta property="og:title" content="Keras documentation: Making new layers and models via subclassing"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Making new layers and models via subclassing"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Making new layers and models via subclassing</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link active" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-sublink" href="/guides/functional_api/">The Functional API</a> <a class="nav-sublink" href="/guides/sequential_model/">The Sequential model</a> <a class="nav-sublink active" href="/guides/making_new_layers_and_models_via_subclassing/">Making new layers & models via subclassing</a> <a class="nav-sublink" href="/guides/training_with_built_in_methods/">Training & evaluation with the built-in methods</a> <a class="nav-sublink" href="/guides/custom_train_step_in_jax/">Customizing `fit()` with JAX</a> <a class="nav-sublink" href="/guides/custom_train_step_in_tensorflow/">Customizing `fit()` with TensorFlow</a> <a class="nav-sublink" href="/guides/custom_train_step_in_torch/">Customizing `fit()` with PyTorch</a> <a class="nav-sublink" href="/guides/writing_a_custom_training_loop_in_jax/">Writing a custom training loop in JAX</a> <a class="nav-sublink" href="/guides/writing_a_custom_training_loop_in_tensorflow/">Writing a custom training loop in TensorFlow</a> <a class="nav-sublink" href="/guides/writing_a_custom_training_loop_in_torch/">Writing a custom training loop in PyTorch</a> <a class="nav-sublink" href="/guides/serialization_and_saving/">Serialization & saving</a> <a class="nav-sublink" href="/guides/customizing_saving_and_serialization/">Customizing saving & serialization</a> <a class="nav-sublink" href="/guides/writing_your_own_callbacks/">Writing your own callbacks</a> <a class="nav-sublink" href="/guides/transfer_learning/">Transfer learning & fine-tuning</a> <a class="nav-sublink" href="/guides/distributed_training_with_jax/">Distributed training with JAX</a> <a class="nav-sublink" href="/guides/distributed_training_with_tensorflow/">Distributed training with TensorFlow</a> <a class="nav-sublink" href="/guides/distributed_training_with_torch/">Distributed training with PyTorch</a> <a class="nav-sublink" href="/guides/distribution/">Distributed training with Keras 3</a> <a class="nav-sublink" href="/guides/migrating_to_keras_3/">Migrating Keras 2 code to Keras 3</a> <a class="nav-link" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparam Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/guides/'>Developer guides</a> / Making new layers and models via subclassing </div> <div class='k-content'> <h1 id="making-new-layers-and-models-via-subclassing">Making new layers and models via subclassing</h1> Author: <a href="https://twitter.com/fchollet">fchollet</a> Date created: 2019/03/01 Last modified: 2023/06/25 Description: Complete guide to writing <code>Layer</code> and <code>Model</code> objects from scratch. <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/guides/ipynb/making_new_layers_and_models_via_subclassing.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/guides/making_new_layers_and_models_via_subclassing.py">GitHub source</a> <hr /> <h2 id="introduction">Introduction</h2> This guide will cover everything you need to know to build your own subclassed layers and models. In particular, you'll learn about the following features: <ul> <li>The <code>Layer</code> class</li> <li>The <code>add_weight()</code> method</li> <li>Trainable and non-trainable weights</li> <li>The <code>build()</code> method</li> <li>Making sure your layers can be used with any backend</li> <li>The <code>add_loss()</code> method</li> <li>The <code>training</code> argument in <code>call()</code></li> <li>The <code>mask</code> argument in <code>call()</code></li> <li>Making sure your layers can be serialized</li> </ul> Let's dive in. <hr /> <h2 id="setup">Setup</h2> <div class="codehilite"><pre><code>import numpy as np import keras from keras import ops from keras import layers </code></pre></div> <hr /> <h2 id="the-layer-class-the-combination-of-state-weights-and-some-computation">The <code>Layer</code> class: the combination of state (weights) and some computation</h2> One of the central abstractions in Keras is the <code>Layer</code> class. A layer encapsulates both a state (the layer's "weights") and a transformation from inputs to outputs (a "call", the layer's forward pass). Here's a densely-connected layer. It has two state variables: the variables <code>w</code> and <code>b</code>. <div class="codehilite"><pre><code>class Linear(keras.layers.Layer): def __init__(self, units=32, input_dim=32): super().__init__() self.w = self.add_weight( shape=(input_dim, units), initializer="random_normal", trainable=True, ) self.b = self.add_weight(shape=(units,), initializer="zeros", trainable=True) def call(self, inputs): return ops.matmul(inputs, self.w) + self.b </code></pre></div> You would use a layer by calling it on some tensor input(s), much like a Python function. <div class="codehilite"><pre><code>x = ops.ones((2, 2)) linear_layer = Linear(4, 2) y = linear_layer(x) print(y) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>[[ 0.085416 -0.06821361 -0.00741937 -0.03429271] [ 0.085416 -0.06821361 -0.00741937 -0.03429271]] </code></pre></div> </div> Note that the weights <code>w</code> and <code>b</code> are automatically tracked by the layer upon being set as layer attributes: <div class="codehilite"><pre><code>assert linear_layer.weights == [linear_layer.w, linear_layer.b] </code></pre></div> <hr /> <h2 id="layers-can-have-nontrainable-weights">Layers can have non-trainable weights</h2> Besides trainable weights, you can add non-trainable weights to a layer as well. Such weights are meant not to be taken into account during backpropagation, when you are training the layer. Here's how to add and use a non-trainable weight: <div class="codehilite"><pre><code>class ComputeSum(keras.layers.Layer): def __init__(self, input_dim): super().__init__() self.total = self.add_weight( initializer="zeros", shape=(input_dim,), trainable=False ) def call(self, inputs): self.total.assign_add(ops.sum(inputs, axis=0)) return self.total x = ops.ones((2, 2)) my_sum = ComputeSum(2) y = my_sum(x) print(y.numpy()) y = my_sum(x) print(y.numpy()) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>[2. 2.] [4. 4.] </code></pre></div> </div> It's part of <code>layer.weights</code>, but it gets categorized as a non-trainable weight: <div class="codehilite"><pre><code>print("weights:", len(my_sum.weights)) print("non-trainable weights:", len(my_sum.non_trainable_weights)) # It's not included in the trainable weights: print("trainable_weights:", my_sum.trainable_weights) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>weights: 1 non-trainable weights: 1 trainable_weights: [] </code></pre></div> </div> <hr /> <h2 id="best-practice-deferring-weight-creation-until-the-shape-of-the-inputs-is-known">Best practice: deferring weight creation until the shape of the inputs is known</h2> Our <code>Linear</code> layer above took an <code>input_dim</code> argument that was used to compute the shape of the weights <code>w</code> and <code>b</code> in <code>__init__()</code>: <div class="codehilite"><pre><code>class Linear(keras.layers.Layer): def __init__(self, units=32, input_dim=32): super().__init__() self.w = self.add_weight( shape=(input_dim, units), initializer="random_normal", trainable=True, ) self.b = self.add_weight(shape=(units,), initializer="zeros", trainable=True) def call(self, inputs): return ops.matmul(inputs, self.w) + self.b </code></pre></div> In many cases, you may not know in advance the size of your inputs, and you would like to lazily create weights when that value becomes known, some time after instantiating the layer. In the Keras API, we recommend creating layer weights in the <code>build(self, inputs_shape)</code> method of your layer. Like this: <div class="codehilite"><pre><code>class Linear(keras.layers.Layer): def __init__(self, units=32): 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="random_normal", trainable=True ) def call(self, inputs): return ops.matmul(inputs, self.w) + self.b </code></pre></div> The <code>__call__()</code> method of your layer will automatically run build the first time it is called. You now have a layer that's lazy and thus easier to use: <div class="codehilite"><pre><code># At instantiation, we don't know on what inputs this is going to get called linear_layer = Linear(32) # The layer's weights are created dynamically the first time the layer is called y = linear_layer(x) </code></pre></div> Implementing <code>build()</code> separately as shown above nicely separates creating weights only once from using weights in every call. <hr /> <h2 id="layers-are-recursively-composable">Layers are recursively composable</h2> If you assign a Layer instance as an attribute of another Layer, the outer layer will start tracking the weights created by the inner layer. We recommend creating such sublayers in the <code>__init__()</code> method and leave it to the first <code>__call__()</code> to trigger building their weights. <div class="codehilite"><pre><code>class MLPBlock(keras.layers.Layer): def __init__(self): super().__init__() self.linear_1 = Linear(32) self.linear_2 = Linear(32) self.linear_3 = Linear(1) def call(self, inputs): x = self.linear_1(inputs) x = keras.activations.relu(x) x = self.linear_2(x) x = keras.activations.relu(x) return self.linear_3(x) mlp = MLPBlock() y = mlp(ops.ones(shape=(3, 64))) # The first call to the `mlp` will create the weights print("weights:", len(mlp.weights)) print("trainable weights:", len(mlp.trainable_weights)) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>weights: 6 trainable weights: 6 </code></pre></div> </div> <hr /> <h2 id="backendagnostic-layers-and-backendspecific-layers">Backend-agnostic layers and backend-specific layers</h2> As long as a layer only uses APIs from the <code>keras.ops</code> namespace (or other Keras namespaces such as <code>keras.activations</code>, <code>keras.random</code>, or <code>keras.layers</code>), then it can be used with any backend – TensorFlow, JAX, or PyTorch. All layers you've seen so far in this guide work with all Keras backends. The <code>keras.ops</code> namespace gives you access to: <ul> <li>The NumPy API, e.g. <code>ops.matmul</code>, <code>ops.sum</code>, <code>ops.reshape</code>, <code>ops.stack</code>, etc.</li> <li>Neural networks-specific APIs such as <code>ops.softmax</code>, <code>ops</code>.conv<code>,</code>ops.binary_crossentropy<code>,</code>ops.relu`, etc.</li> </ul> You can also use backend-native APIs in your layers (such as <a href="https://www.tensorflow.org/api_docs/python/tf/nn"><code>tf.nn</code></a> functions), but if you do this, then your layer will only be usable with the backend in question. For instance, you could write the following JAX-specific layer using <code>jax.numpy</code>: <div class="codehilite"><pre><code>import jax class Linear(keras.layers.Layer): ... def call(self, inputs): return jax.numpy.matmul(inputs, self.w) + self.b </code></pre></div> This would be the equivalent TensorFlow-specific layer: <div class="codehilite"><pre><code>import tensorflow as tf class Linear(keras.layers.Layer): ... def call(self, inputs): return tf.matmul(inputs, self.w) + self.b </code></pre></div> And this would be the equivalent PyTorch-specific layer: <div class="codehilite"><pre><code>import torch class Linear(keras.layers.Layer): ... def call(self, inputs): return torch.matmul(inputs, self.w) + self.b </code></pre></div> Because cross-backend compatibility is a tremendously useful property, we strongly recommend that you seek to always make your layers backend-agnostic by leveraging only Keras APIs. <hr /> <h2 id="the-addloss-method">The <code>add_loss()</code> method</h2> When writing the <code>call()</code> method of a layer, you can create loss tensors that you will want to use later, when writing your training loop. This is doable by calling <code>self.add_loss(value)</code>: <div class="codehilite"><pre><code># A layer that creates an activity regularization loss class ActivityRegularizationLayer(keras.layers.Layer): def __init__(self, rate=1e-2): super().__init__() self.rate = rate def call(self, inputs): self.add_loss(self.rate * ops.mean(inputs)) return inputs </code></pre></div> These losses (including those created by any inner layer) can be retrieved via <code>layer.losses</code>. This property is reset at the start of every <code>__call__()</code> to the top-level layer, so that <code>layer.losses</code> always contains the loss values created during the last forward pass. <div class="codehilite"><pre><code>class OuterLayer(keras.layers.Layer): def __init__(self): super().__init__() self.activity_reg = ActivityRegularizationLayer(1e-2) def call(self, inputs): return self.activity_reg(inputs) layer = OuterLayer() assert len(layer.losses) == 0 # No losses yet since the layer has never been called _ = layer(ops.zeros((1, 1))) assert len(layer.losses) == 1 # We created one loss value # `layer.losses` gets reset at the start of each __call__ _ = layer(ops.zeros((1, 1))) assert len(layer.losses) == 1 # This is the loss created during the call above </code></pre></div> In addition, the <code>loss</code> property also contains regularization losses created for the weights of any inner layer: <div class="codehilite"><pre><code>class OuterLayerWithKernelRegularizer(keras.layers.Layer): def __init__(self): super().__init__() self.dense = keras.layers.Dense( 32, kernel_regularizer=keras.regularizers.l2(1e-3) ) def call(self, inputs): return self.dense(inputs) layer = OuterLayerWithKernelRegularizer() _ = layer(ops.zeros((1, 1))) # This is `1e-3 * sum(layer.dense.kernel ** 2)`, # created by the `kernel_regularizer` above. print(layer.losses) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>[Array(0.00217911, dtype=float32)] </code></pre></div> </div> These losses are meant to be taken into account when writing custom training loops. They also work seamlessly with <code>fit()</code> (they get automatically summed and added to the main loss, if any): <div class="codehilite"><pre><code>inputs = keras.Input(shape=(3,)) outputs = ActivityRegularizationLayer()(inputs) model = keras.Model(inputs, outputs) # If there is a loss passed in `compile`, the regularization # losses get added to it model.compile(optimizer="adam", loss="mse") model.fit(np.random.random((2, 3)), np.random.random((2, 3))) # It's also possible not to pass any loss in `compile`, # since the model already has a loss to minimize, via the `add_loss` # call during the forward pass! model.compile(optimizer="adam") model.fit(np.random.random((2, 3)), np.random.random((2, 3))) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 60ms/step - loss: 0.2650 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 15ms/step - loss: 0.0050 <keras.src.callbacks.history.History at 0x146f71960> </code></pre></div> </div> <hr /> <h2 id="you-can-optionally-enable-serialization-on-your-layers">You can optionally enable serialization on your layers</h2> If you need your custom layers to be serializable as part of a <a href="/guides/functional_api/">Functional model</a>, you can optionally implement a <code>get_config()</code> method: <div class="codehilite"><pre><code>class Linear(keras.layers.Layer): def __init__(self, units=32): 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="random_normal", trainable=True ) def call(self, inputs): return ops.matmul(inputs, self.w) + self.b def get_config(self): return {"units": self.units} # Now you can recreate the layer from its config: layer = Linear(64) config = layer.get_config() print(config) new_layer = Linear.from_config(config) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>{'units': 64} </code></pre></div> </div> Note that the <code>__init__()</code> method of the base <code>Layer</code> class takes some keyword arguments, in particular a <code>name</code> and a <code>dtype</code>. It's good practice to pass these arguments to the parent class in <code>__init__()</code> and to include them in the layer config: <div class="codehilite"><pre><code>class Linear(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="random_normal", trainable=True, ) self.b = self.add_weight( shape=(self.units,), initializer="random_normal", trainable=True ) def call(self, inputs): return ops.matmul(inputs, self.w) + self.b def get_config(self): config = super().get_config() config.update({"units": self.units}) return config layer = Linear(64) config = layer.get_config() print(config) new_layer = Linear.from_config(config) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>{'name': 'linear_7', 'trainable': True, 'dtype': 'float32', 'units': 64} </code></pre></div> </div> If you need more flexibility when deserializing the layer from its config, you can also override the <code>from_config()</code> class method. This is the base implementation of <code>from_config()</code>: <div class="codehilite"><pre><code>def from_config(cls, config): return cls(**config) </code></pre></div> To learn more about serialization and saving, see the complete <a href="/guides/serialization_and_saving/">guide to saving and serializing models</a>. <hr /> <h2 id="privileged-training-argument-in-the-call-method">Privileged <code>training</code> argument in the <code>call()</code> method</h2> Some layers, in particular the <code>BatchNormalization</code> layer and the <code>Dropout</code> layer, have different behaviors during training and inference. For such layers, it is standard practice to expose a <code>training</code> (boolean) argument in the <code>call()</code> method. By exposing this argument in <code>call()</code>, you enable the built-in training and evaluation loops (e.g. <code>fit()</code>) to correctly use the layer in training and inference. <div class="codehilite"><pre><code>class CustomDropout(keras.layers.Layer): def __init__(self, rate, **kwargs): super().__init__(**kwargs) self.rate = rate self.seed_generator = keras.random.SeedGenerator(1337) def call(self, inputs, training=None): if training: return keras.random.dropout( inputs, rate=self.rate, seed=self.seed_generator ) return inputs </code></pre></div> <hr /> <h2 id="privileged-mask-argument-in-the-call-method">Privileged <code>mask</code> argument in the <code>call()</code> method</h2> The other privileged argument supported by <code>call()</code> is the <code>mask</code> argument. You will find it in all Keras RNN layers. A mask is a boolean tensor (one boolean value per timestep in the input) used to skip certain input timesteps when processing timeseries data. Keras will automatically pass the correct <code>mask</code> argument to <code>__call__()</code> for layers that support it, when a mask is generated by a prior layer. Mask-generating layers are the <code>Embedding</code> layer configured with <code>mask_zero=True</code>, and the <code>Masking</code> layer. <hr /> <h2 id="the-model-class">The <code>Model</code> class</h2> In general, you will use the <code>Layer</code> class to define inner computation blocks, and will use the <code>Model</code> class to define the outer model – the object you will train. For instance, in a ResNet50 model, you would have several ResNet blocks subclassing <code>Layer</code>, and a single <code>Model</code> encompassing the entire ResNet50 network. The <code>Model</code> class has the same API as <code>Layer</code>, with the following differences: <ul> <li>It exposes built-in training, evaluation, and prediction loops (<code>model.fit()</code>, <code>model.evaluate()</code>, <code>model.predict()</code>).</li> <li>It exposes the list of its inner layers, via the <code>model.layers</code> property.</li> <li>It exposes saving and serialization APIs (<code>save()</code>, <code>save_weights()</code>...)</li> </ul> Effectively, the <code>Layer</code> class corresponds to what we refer to in the literature as a "layer" (as in "convolution layer" or "recurrent layer") or as a "block" (as in "ResNet block" or "Inception block"). Meanwhile, the <code>Model</code> class corresponds to what is referred to in the literature as a "model" (as in "deep learning model") or as a "network" (as in "deep neural network"). So if you're wondering, "should I use the <code>Layer</code> class or the <code>Model</code> class?", ask yourself: will I need to call <code>fit()</code> on it? Will I need to call <code>save()</code> on it? If so, go with <code>Model</code>. If not (either because your class is just a block in a bigger system, or because you are writing training & saving code yourself), use <code>Layer</code>. For instance, we could take our mini-resnet example above, and use it to build a <code>Model</code> that we could train with <code>fit()</code>, and that we could save with <code>save_weights()</code>: <div class="codehilite"><pre><code>class ResNet(keras.Model): def __init__(self, num_classes=1000): super().__init__() self.block_1 = ResNetBlock() self.block_2 = ResNetBlock() self.global_pool = layers.GlobalAveragePooling2D() self.classifier = Dense(num_classes) def call(self, inputs): x = self.block_1(inputs) x = self.block_2(x) x = self.global_pool(x) return self.classifier(x) resnet = ResNet() dataset = ... resnet.fit(dataset, epochs=10) resnet.save(filepath.keras) </code></pre></div> <hr /> <h2 id="putting-it-all-together-an-endtoend-example">Putting it all together: an end-to-end example</h2> Here's what you've learned so far: <ul> <li>A <code>Layer</code> encapsulate a state (created in <code>__init__()</code> or <code>build()</code>) and some computation (defined in <code>call()</code>).</li> <li>Layers can be recursively nested to create new, bigger computation blocks.</li> <li>Layers are backend-agnostic as long as they only use Keras APIs. You can use backend-native APIs (such as <code>jax.numpy</code>, <code>torch.nn</code> or <a href="https://www.tensorflow.org/api_docs/python/tf/nn"><code>tf.nn</code></a>), but then your layer will only be usable with that specific backend.</li> <li>Layers can create and track losses (typically regularization losses) via <code>add_loss()</code>.</li> <li>The outer container, the thing you want to train, is a <code>Model</code>. A <code>Model</code> is just like a <code>Layer</code>, but with added training and serialization utilities.</li> </ul> Let's put all of these things together into an end-to-end example: we're going to implement a Variational AutoEncoder (VAE) in a backend-agnostic fashion – so that it runs the same with TensorFlow, JAX, and PyTorch. We'll train it on MNIST digits. Our VAE will be a subclass of <code>Model</code>, built as a nested composition of layers that subclass <code>Layer</code>. It will feature a regularization loss (KL divergence). <div class="codehilite"><pre><code>class Sampling(layers.Layer): """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit.""" def __init__(self, **kwargs): super().__init__(**kwargs) self.seed_generator = keras.random.SeedGenerator(1337) def call(self, inputs): z_mean, z_log_var = inputs batch = ops.shape(z_mean)[0] dim = ops.shape(z_mean)[1] epsilon = keras.random.normal(shape=(batch, dim), seed=self.seed_generator) return z_mean + ops.exp(0.5 * z_log_var) * epsilon class Encoder(layers.Layer): """Maps MNIST digits to a triplet (z_mean, z_log_var, z).""" def __init__(self, latent_dim=32, intermediate_dim=64, name="encoder", **kwargs): super().__init__(name=name, **kwargs) self.dense_proj = layers.Dense(intermediate_dim, activation="relu") self.dense_mean = layers.Dense(latent_dim) self.dense_log_var = layers.Dense(latent_dim) self.sampling = Sampling() def call(self, inputs): x = self.dense_proj(inputs) z_mean = self.dense_mean(x) z_log_var = self.dense_log_var(x) z = self.sampling((z_mean, z_log_var)) return z_mean, z_log_var, z class Decoder(layers.Layer): """Converts z, the encoded digit vector, back into a readable digit.""" def __init__(self, original_dim, intermediate_dim=64, name="decoder", **kwargs): super().__init__(name=name, **kwargs) self.dense_proj = layers.Dense(intermediate_dim, activation="relu") self.dense_output = layers.Dense(original_dim, activation="sigmoid") def call(self, inputs): x = self.dense_proj(inputs) return self.dense_output(x) class VariationalAutoEncoder(keras.Model): """Combines the encoder and decoder into an end-to-end model for training.""" def __init__( self, original_dim, intermediate_dim=64, latent_dim=32, name="autoencoder", **kwargs ): super().__init__(name=name, **kwargs) self.original_dim = original_dim self.encoder = Encoder(latent_dim=latent_dim, intermediate_dim=intermediate_dim) self.decoder = Decoder(original_dim, intermediate_dim=intermediate_dim) def call(self, inputs): z_mean, z_log_var, z = self.encoder(inputs) reconstructed = self.decoder(z) # Add KL divergence regularization loss. kl_loss = -0.5 * ops.mean( z_log_var - ops.square(z_mean) - ops.exp(z_log_var) + 1 ) self.add_loss(kl_loss) return reconstructed </code></pre></div> Let's train it on MNIST using the <code>fit()</code> API: <div class="codehilite"><pre><code>(x_train, _), _ = keras.datasets.mnist.load_data() x_train = x_train.reshape(60000, 784).astype("float32") / 255 original_dim = 784 vae = VariationalAutoEncoder(784, 64, 32) optimizer = keras.optimizers.Adam(learning_rate=1e-3) vae.compile(optimizer, loss=keras.losses.MeanSquaredError()) vae.fit(x_train, x_train, epochs=2, batch_size=64) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Epoch 1/2 938/938 ━━━━━━━━━━━━━━━━━━━━ 2s 1ms/step - loss: 0.0942 Epoch 2/2 938/938 ━━━━━━━━━━━━━━━━━━━━ 1s 859us/step - loss: 0.0677 <keras.src.callbacks.history.History at 0x146fe62f0> </code></pre></div> </div> </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#making-new-layers-and-models-via-subclassing'>Making new layers and models via subclassing</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#setup'>Setup</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#the-layer-class-the-combination-of-state-weights-and-some-computation'>The <code>Layer</code> class: the combination of state (weights) and some computation</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#layers-can-have-nontrainable-weights'>Layers can have non-trainable weights</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#best-practice-deferring-weight-creation-until-the-shape-of-the-inputs-is-known'>Best practice: deferring weight creation until the shape of the inputs is known</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#layers-are-recursively-composable'>Layers are recursively composable</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#backendagnostic-layers-and-backendspecific-layers'>Backend-agnostic layers and backend-specific layers</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#the-addloss-method'>The <code>add_loss()</code> method</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#you-can-optionally-enable-serialization-on-your-layers'>You can optionally enable serialization on your layers</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#privileged-training-argument-in-the-call-method'>Privileged <code>training</code> argument in the <code>call()</code> method</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#privileged-mask-argument-in-the-call-method'>Privileged <code>mask</code> argument in the <code>call()</code> method</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#the-model-class'>The <code>Model</code> class</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#putting-it-all-together-an-endtoend-example'>Putting it all together: an end-to-end example</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>