How to save the architecture of a tensorflow model (summary) in a json file ?

Example of how to show save the architecture of a tensorflow model (summary) in a json file:

Create a model

Let's create a simple untrained model with TensorFlow:

from keras.utils.data_utils import get_file
from tensorflow import keras
from tensorflow.keras import layers

model = keras.Sequential([
    layers.Dense(20, activation='relu', input_shape=[11]),
    layers.Dense(10, activation='relu'),
    layers.Dense(10, activation='relu'),
    layers.Dense(1, activation='sigmoid')
    ])

Get model summary

To get a summary of the model, a solution is to use Module: tf.summary:

model.summary()

returns

Model: "sequential"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 dense (Dense)               (None, 20)                240

 dense_1 (Dense)             (None, 10)                210

 dense_2 (Dense)             (None, 10)                110

 dense_3 (Dense)             (None, 1)                 11

=================================================================
Total params: 571
Trainable params: 571
Non-trainable params: 0
_________________________________________________________________

Save model architecture in a json file

To save a model architecture in a json file, a solution is to use to_json() (see Save and load Keras models):

json_config = model.to_json()

returns here

'{"class_name": "Sequential", "config": {"name": "sequential", "layers": [{"class_name": "InputLayer", "config": {"batch_input_shape": [null, 11], "dtype": "float32", "sparse": false, "ragged": false, "name": "dense_input"}}, {"class_name": "Dense", "config": {"name": "dense", "trainable": true, "batch_input_shape": [null, 11], "dtype": "float32", "units": 20, "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "GlorotUniform", "config": {"seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}}, {"class_name": "Dense", "config": {"name": "dense_1", "trainable": true, "dtype": "float32", "units": 10, "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "GlorotUniform", "config": {"seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}}, {"class_name": "Dense", "config": {"name": "dense_2", "trainable": true, "dtype": "float32", "units": 10, "activation": "relu", "use_bias": true, "kernel_initializer": {"class_name": "GlorotUniform", "config": {"seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}}, {"class_name": "Dense", "config": {"name": "dense_3", "trainable": true, "dtype": "float32", "units": 1, "activation": "sigmoid", "use_bias": true, "kernel_initializer": {"class_name": "GlorotUniform", "config": {"seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}}]}, "keras_version": "2.7.0", "backend": "tensorflow"}'

and save the model architecture in a json file called for example: model_architecture.json:

import json

with open('model_architecture.json', 'w') as fp:
        json.dump(json_config, fp)

Reload model architecture

Now, let's try to read the model architecture saved in the json file:

import json

with open('model_architecture.json') as json_data:
    print(type(json_data))
    json_config = json.load(json_data)

new_model = keras.models.model_from_json(json_config)

new_model.summary()

returns

<class '_io.TextIOWrapper'>
Model: "sequential"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 dense (Dense)               (None, 20)                240

 dense_1 (Dense)             (None, 10)                210

 dense_2 (Dense)             (None, 10)                110

 dense_3 (Dense)             (None, 1)                 11

=================================================================
Total params: 571
Trainable params: 571
Non-trainable params: 0
_________________________________________________________________