Compatibility between tensorboard and keras3/keras-core to display functional model graphs

[nhuet]

When using keras-nightly which is now a preview of the keras 3.0 to come, tensorboard fails to represent the conceptual graph of a functional keras model.

I know that tensorboard is intended to work with the keras of same version, so i put this as a feature request instead of a bug. I tried also with tensorboard-nightly with the same result: sequential models are well represented but not functional models, probably because of a change in json representation. The error comes from "inbound_nodes" which can now be a list of dictionaries/lists instead of a list of lists.

Here is a minimal example showing the bug:

import json

from keras import Input, Model, Sequential
from keras.layers import Dense

from tensorboard.plugins.graph.keras_util import keras_model_to_graph_def


input = Input((1,))
layer = Dense(3)

## sequential: working
model_seq = Sequential(layers=[input, layer])
graph_def = keras_model_to_graph_def(json.loads(model_seq.to_json()))

## functional: not working  (with keras-nightly aka keras3)
output = layer(input)
model_func = Model(input, output)
graph_def = keras_model_to_graph_def(json.loads(model_func.to_json()))

which results in

Traceback (most recent call last):
File "/home/nolwen/Projects/decomon/tensorboard/minimal_example.py", line 19, in
graph_def = keras_model_to_graph_def(json.loads(model_func.to_json()))
File "/home/nolwen/Projects/decomon/tensorboard/tb-keras3-venv/lib/python3.9/site-packages/tensorboard/plugins/graph/keras_util.py", line 248, in keras_model_to_graph_def
inbound_nodes = _norm_to_list_of_layers(maybe_inbound_node)
File "/home/nolwen/Projects/decomon/tensorboard/tb-keras3-venv/lib/python3.9/site-packages/tensorboard/plugins/graph/keras_util.py", line 116, in _norm_to_list_of_layers
maybe_layers if isinstance(maybe_layers[0], (list,)) else [maybe_layers]
KeyError: 0

Here is an extract of the json reprenstation showing that inbound_nodes members are not always lists:

{

    "inbound_nodes": [
          {
            "args": [
              {
                "class_name": "__keras_tensor__",
                "config": {
                  "shape": [
                    null,
                    1
                  ],
                  "dtype": "float32",
                  "keras_history": [
                    "input_layer",
                    0,
                    0
                  ]
                }
              }
            ],
            "kwargs": {}
          }
        ]
}

You have to run it in an environment with tensorboard (potentially tb-nightly) and keras-nightly. For instance with:

python -m venv tb-keras3-venv
. tb-keras3-venv/bin/activate
pip install tf-nightly
pip uninstall tf-keras-nightly keras-nightly -y  # be sure to avoid a mix of both keras
pip install keras-nightly

[JamesHollyer]

Thank you for this detailed issue! I have added this to out internal hotlist. I hope we can get to it soon!

[yatbear]

More context: we recently pinned to Keras 2 before fully migrating our codebase to Keras 3: #6684

Apologies for the inconvenience!

[nhuet]

Ok, i understand, thanks for the explanation.

[nhuet]

Hi, is there any progress on this issue? Since 2.16, tensorflow and tensorboard are supposed to be keras 3 fully compatible since this is now the default keras installed. It seems that functional keras models cannot still be displayed in graph plugin. (Same error, inbound_nodes list is made of dicts instead of a lists, thus missing the key 0.

[yatbear]

Hi @nhuet,

Apologies since we haven't got any bandwidth to work on this issue, will post an update here as soon as there's any update.

[yatbear]

Hi @nhuet,

Keras 3 compatibility issue for graph plugin should be fixed at #6823 and #6826. This change will be available in tb-nightly in the next one or two days, please try again and let us know if this resolves your issue. Thanks for your patience!

[miticollo]

Hi @yatbear!

Thanks for these two commits. But I found another issue.

Look this notebook, it seems that graph plugin fails when you want to visualize the conceptual graph. In particular for ResNet-152 an error is risen here.
If you look the output of model.to_json() you can notice the following JSON list:

[
    {
        "class_name": "__ellipsis__",
        "config": {}
    },
    2
]

In this case, arg is 2 and is an int. So the method .get("config", {}) doesn't exist. The int type is compatible with serialization like shown here. Only KerasTensor has _keras_history.