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/customizing_saving_and_serialization/" />  <meta property="og:title" content="Keras documentation: Customizing Saving and Serialization"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Customizing Saving and Serialization"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Customizing Saving and Serialization</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" 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 active" 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-sublink" href="/guides/keras_tuner/">Hyperparameter Tuning</a> <a class="nav-sublink" href="/guides/keras_cv/">KerasCV</a> <a class="nav-sublink" href="/guides/keras_nlp/">KerasNLP</a> <a class="nav-sublink" href="/guides/keras_hub/">KerasHub</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="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparameter Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> <a class="nav-link" href="/keras_cv/" role="tab" aria-selected="">KerasCV: Computer Vision Workflows</a> <a class="nav-link" href="/keras_nlp/" role="tab" aria-selected="">KerasNLP: Natural Language Workflows</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> / Customizing Saving and Serialization </div> <div class='k-content'> <h1 id="customizing-saving-and-serialization">Customizing Saving and Serialization</h1> Author: Neel Kovelamudi Date created: 2023/03/15 Last modified: 2023/03/15 Description: A more advanced guide on customizing saving for your layers and models. <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/customizing_saving_and_serialization.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/customizing_saving_and_serialization.py">GitHub source</a> <hr /> <h2 id="introduction">Introduction</h2> This guide covers advanced methods that can be customized in Keras saving. For most users, the methods outlined in the primary <a href="https://keras.io/guides/serialization_and_saving">Serialize, save, and export guide</a> are sufficient. <h3 id="apis">APIs</h3> We will cover the following APIs: <ul> <li><code>save_assets()</code> and <code>load_assets()</code></li> <li><code>save_own_variables()</code> and <code>load_own_variables()</code></li> <li><code>get_build_config()</code> and <code>build_from_config()</code></li> <li><code>get_compile_config()</code> and <code>compile_from_config()</code></li> </ul> When restoring a model, these get executed in the following order: <ul> <li><code>build_from_config()</code></li> <li><code>compile_from_config()</code></li> <li><code>load_own_variables()</code></li> <li><code>load_assets()</code></li> </ul> <hr /> <h2 id="setup">Setup</h2> <div class="codehilite"><pre><code>import os import numpy as np import keras </code></pre></div> <hr /> <h2 id="state-saving-customization">State saving customization</h2> These methods determine how the state of your model's layers is saved when calling <code>model.save()</code>. You can override them to take full control of the state saving process. <h3 id="saveownvariables-and-loadownvariables"><code>save_own_variables()</code> and <code>load_own_variables()</code></h3> These methods save and load the state variables of the layer when <code>model.save()</code> and <code>keras.models.load_model()</code> are called, respectively. By default, the state variables saved and loaded are the weights of the layer (both trainable and non-trainable). Here is the default implementation of <code>save_own_variables()</code>: <div class="codehilite"><pre><code>def save_own_variables(self, store): all_vars = self._trainable_weights + self._non_trainable_weights for i, v in enumerate(all_vars): store[f"{i}"] = v.numpy() </code></pre></div> The store used by these methods is a dictionary that can be populated with the layer variables. Let's take a look at an example customizing this. Example: <div class="codehilite"><pre><code>@keras.utils.register_keras_serializable(package="my_custom_package") class LayerWithCustomVariable(keras.layers.Dense): def __init__(self, units, **kwargs): super().__init__(units, **kwargs) self.my_variable = keras.Variable( np.random.random((units,)), name="my_variable", dtype="float32" ) def save_own_variables(self, store): super().save_own_variables(store) # Stores the value of the variable upon saving store["variables"] = self.my_variable.numpy() def load_own_variables(self, store): # Assigns the value of the variable upon loading self.my_variable.assign(store["variables"]) # Load the remaining weights for i, v in enumerate(self.weights): v.assign(store[f"{i}"]) # Note: You must specify how all variables (including layer weights) # are loaded in `load_own_variables.` def call(self, inputs): dense_out = super().call(inputs) return dense_out + self.my_variable model = keras.Sequential([LayerWithCustomVariable(1)]) ref_input = np.random.random((8, 10)) ref_output = np.random.random((8, 10)) model.compile(optimizer="adam", loss="mean_squared_error") model.fit(ref_input, ref_output) model.save("custom_vars_model.keras") restored_model = keras.models.load_model("custom_vars_model.keras") np.testing.assert_allclose( model.layers[0].my_variable.numpy(), restored_model.layers[0].my_variable.numpy(), ) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 101ms/step - loss: 0.2908 </code></pre></div> </div> <h3 id="saveassets-and-loadassets"><code>save_assets()</code> and <code>load_assets()</code></h3> These methods can be added to your model class definition to store and load any additional information that your model needs. For example, NLP domain layers such as TextVectorization layers and IndexLookup layers may need to store their associated vocabulary (or lookup table) in a text file upon saving. Let's take at the basics of this workflow with a simple file <code>assets.txt</code>. Example: <div class="codehilite"><pre><code>@keras.saving.register_keras_serializable(package="my_custom_package") class LayerWithCustomAssets(keras.layers.Dense): def __init__(self, vocab=None, *args, **kwargs): super().__init__(*args, **kwargs) self.vocab = vocab def save_assets(self, inner_path): # Writes the vocab (sentence) to text file at save time. with open(os.path.join(inner_path, "vocabulary.txt"), "w") as f: f.write(self.vocab) def load_assets(self, inner_path): # Reads the vocab (sentence) from text file at load time. with open(os.path.join(inner_path, "vocabulary.txt"), "r") as f: text = f.read() self.vocab = text.replace("<unk>", "little") model = keras.Sequential( [LayerWithCustomAssets(vocab="Mary had a <unk> lamb.", units=5)] ) x = np.random.random((10, 10)) y = model(x) model.save("custom_assets_model.keras") restored_model = keras.models.load_model("custom_assets_model.keras") np.testing.assert_string_equal( restored_model.layers[0].vocab, "Mary had a little lamb." ) </code></pre></div> <hr /> <h2 id="build-and-compile-saving-customization"><code>build</code> and <code>compile</code> saving customization</h2> <h3 id="getbuildconfig-and-buildfromconfig"><code>get_build_config()</code> and <code>build_from_config()</code></h3> These methods work together to save the layer's built states and restore them upon loading. By default, this only includes a build config dictionary with the layer's input shape, but overriding these methods can be used to include further Variables and Lookup Tables that can be useful to restore for your built model. Example: <div class="codehilite"><pre><code>@keras.saving.register_keras_serializable(package="my_custom_package") class LayerWithCustomBuild(keras.layers.Layer): def __init__(self, units=32, **kwargs): super().__init__(**kwargs) self.units = units def call(self, inputs): return keras.ops.matmul(inputs, self.w) + self.b def get_config(self): return dict(units=self.units, **super().get_config()) def build(self, input_shape, layer_init): # Note the overriding of `build()` to add an extra argument. # Therefore, we will need to manually call build with `layer_init` argument # before the first execution of `call()`. super().build(input_shape) self._input_shape = input_shape self.w = self.add_weight( shape=(input_shape[-1], self.units), initializer=layer_init, trainable=True, ) self.b = self.add_weight( shape=(self.units,), initializer=layer_init, trainable=True, ) self.layer_init = layer_init def get_build_config(self): build_config = { "layer_init": self.layer_init, "input_shape": self._input_shape, } # Stores our initializer for `build()` return build_config def build_from_config(self, config): # Calls `build()` with the parameters at loading time self.build(config["input_shape"], config["layer_init"]) custom_layer = LayerWithCustomBuild(units=16) custom_layer.build(input_shape=(8,), layer_init="random_normal") model = keras.Sequential( [ custom_layer, keras.layers.Dense(1, activation="sigmoid"), ] ) x = np.random.random((16, 8)) y = model(x) model.save("custom_build_model.keras") restored_model = keras.models.load_model("custom_build_model.keras") np.testing.assert_equal(restored_model.layers[0].layer_init, "random_normal") np.testing.assert_equal(restored_model.built, True) </code></pre></div> <h3 id="getcompileconfig-and-compilefromconfig"><code>get_compile_config()</code> and <code>compile_from_config()</code></h3> These methods work together to save the information with which the model was compiled (optimizers, losses, etc.) and restore and re-compile the model with this information. Overriding these methods can be useful for compiling the restored model with custom optimizers, custom losses, etc., as these will need to be deserialized prior to calling <code>model.compile</code> in <code>compile_from_config()</code>. Let's take a look at an example of this. Example: <div class="codehilite"><pre><code>@keras.saving.register_keras_serializable(package="my_custom_package") def small_square_sum_loss(y_true, y_pred): loss = keras.ops.square(y_pred - y_true) loss = loss / 10.0 loss = keras.ops.sum(loss, axis=1) return loss @keras.saving.register_keras_serializable(package="my_custom_package") def mean_pred(y_true, y_pred): return keras.ops.mean(y_pred) @keras.saving.register_keras_serializable(package="my_custom_package") class ModelWithCustomCompile(keras.Model): def __init__(self, **kwargs): super().__init__(**kwargs) self.dense1 = keras.layers.Dense(8, activation="relu") self.dense2 = keras.layers.Dense(4, activation="softmax") def call(self, inputs): x = self.dense1(inputs) return self.dense2(x) def compile(self, optimizer, loss_fn, metrics): super().compile(optimizer=optimizer, loss=loss_fn, metrics=metrics) self.model_optimizer = optimizer self.loss_fn = loss_fn self.loss_metrics = metrics def get_compile_config(self): # These parameters will be serialized at saving time. return { "model_optimizer": self.model_optimizer, "loss_fn": self.loss_fn, "metric": self.loss_metrics, } def compile_from_config(self, config): # Deserializes the compile parameters (important, since many are custom) optimizer = keras.utils.deserialize_keras_object(config["model_optimizer"]) loss_fn = keras.utils.deserialize_keras_object(config["loss_fn"]) metrics = keras.utils.deserialize_keras_object(config["metric"]) # Calls compile with the deserialized parameters self.compile(optimizer=optimizer, loss_fn=loss_fn, metrics=metrics) model = ModelWithCustomCompile() model.compile( optimizer="SGD", loss_fn=small_square_sum_loss, metrics=["accuracy", mean_pred] ) x = np.random.random((4, 8)) y = np.random.random((4,)) model.fit(x, y) model.save("custom_compile_model.keras") restored_model = keras.models.load_model("custom_compile_model.keras") np.testing.assert_equal(model.model_optimizer, restored_model.model_optimizer) np.testing.assert_equal(model.loss_fn, restored_model.loss_fn) np.testing.assert_equal(model.loss_metrics, restored_model.loss_metrics) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 79ms/step - accuracy: 0.0000e+00 - loss: 0.0627 - mean_metric_wrapper: 0.2500 </code></pre></div> </div> <hr /> <h2 id="conclusion">Conclusion</h2> Using the methods learned in this tutorial allows for a wide variety of use cases, allowing the saving and loading of complex models with exotic assets and state elements. To recap: <ul> <li><code>save_own_variables</code> and <code>load_own_variables</code> determine how your states are saved and loaded.</li> <li><code>save_assets</code> and <code>load_assets</code> can be added to store and load any additional information your model needs.</li> <li><code>get_build_config</code> and <code>build_from_config</code> save and restore the model's built states.</li> <li><code>get_compile_config</code> and <code>compile_from_config</code> save and restore the model's compiled states.</li> </ul> </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#customizing-saving-and-serialization'>Customizing Saving and Serialization</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-3'> <a href='#apis'>APIs</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#setup'>Setup</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#state-saving-customization'>State saving customization</a> </div> <div class='k-outline-depth-3'> <a href='#saveownvariables-and-loadownvariables'><code>save_own_variables()</code> and <code>load_own_variables()</code></a> </div> <div class='k-outline-depth-3'> <a href='#saveassets-and-loadassets'><code>save_assets()</code> and <code>load_assets()</code></a> </div> <div class='k-outline-depth-2'> ◆ <a href='#build-and-compile-saving-customization'><code>build</code> and <code>compile</code> saving customization</a> </div> <div class='k-outline-depth-3'> <a href='#getbuildconfig-and-buildfromconfig'><code>get_build_config()</code> and <code>build_from_config()</code></a> </div> <div class='k-outline-depth-3'> <a href='#getcompileconfig-and-compilefromconfig'><code>get_compile_config()</code> and <code>compile_from_config()</code></a> </div> <div class='k-outline-depth-2'> ◆ <a href='#conclusion'>Conclusion</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>