Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| # Copyright 2019 The TensorFlow Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ============================================================================== | |
| """Utilities for handling Keras model in graph plugin. | |
| Two canonical types of Keras model are Functional and Sequential. | |
| A model can be serialized as JSON and deserialized to reconstruct a model. | |
| This utility helps with dealing with the serialized Keras model. | |
| They have distinct structures to the configurations in shapes below: | |
| Functional: | |
| config | |
| name: Name of the model. If not specified, it is 'model' with | |
| an optional suffix if there are more than one instance. | |
| input_layers: Keras.layers.Inputs in the model. | |
| output_layers: Layer names that are outputs of the model. | |
| layers: list of layer configurations. | |
| layer: [*] | |
| inbound_nodes: inputs to this layer. | |
| Sequential: | |
| config | |
| name: Name of the model. If not specified, it is 'sequential' with | |
| an optional suffix if there are more than one instance. | |
| layers: list of layer configurations. | |
| layer: [*] | |
| [*]: Note that a model can be a layer. | |
| Please refer to https://github.com/tensorflow/tfjs-layers/blob/master/src/keras_format/model_serialization.ts | |
| for more complete definition. | |
| """ | |
| from tensorboard.compat.proto.graph_pb2 import GraphDef | |
| from tensorboard.compat.tensorflow_stub import dtypes | |
| from tensorboard.util import tb_logging | |
| logger = tb_logging.get_logger() | |
| def _walk_layers(keras_layer): | |
| """Walks the nested keras layer configuration in preorder. | |
| Args: | |
| keras_layer: Keras configuration from model.to_json. | |
| Yields: | |
| A tuple of (name_scope, layer_config). | |
| name_scope: a string representing a scope name, similar to that of tf.name_scope. | |
| layer_config: a dict representing a Keras layer configuration. | |
| """ | |
| yield ("", keras_layer) | |
| if keras_layer.get("config").get("layers"): | |
| name_scope = keras_layer.get("config").get("name") | |
| for layer in keras_layer.get("config").get("layers"): | |
| for sub_name_scope, sublayer in _walk_layers(layer): | |
| sub_name_scope = ( | |
| "%s/%s" % (name_scope, sub_name_scope) | |
| if sub_name_scope | |
| else name_scope | |
| ) | |
| yield (sub_name_scope, sublayer) | |
| def _scoped_name(name_scope, node_name): | |
| """Returns scoped name for a node as a string in the form '<scope>/<node | |
| name>'. | |
| Args: | |
| name_scope: a string representing a scope name, similar to that of tf.name_scope. | |
| node_name: a string representing the current node name. | |
| Returns | |
| A string representing a scoped name. | |
| """ | |
| if name_scope: | |
| return "%s/%s" % (name_scope, node_name) | |
| return node_name | |
| def _is_model(layer): | |
| """Returns True if layer is a model. | |
| Args: | |
| layer: a dict representing a Keras model configuration. | |
| Returns: | |
| bool: True if layer is a model. | |
| """ | |
| return layer.get("config").get("layers") is not None | |
| def _norm_to_list_of_layers(maybe_layers): | |
| """Normalizes to a list of layers. | |
| Args: | |
| maybe_layers: A list of data[1] or a list of list of data. | |
| Returns: | |
| List of list of data. | |
| [1]: A Functional model has fields 'inbound_nodes' and 'output_layers' which can | |
| look like below: | |
| - ['in_layer_name', 0, 0] | |
| - [['in_layer_is_model', 1, 0], ['in_layer_is_model', 1, 1]] | |
| The data inside the list seems to describe [name, size, index]. | |
| """ | |
| return ( | |
| maybe_layers if isinstance(maybe_layers[0], (list,)) else [maybe_layers] | |
| ) | |
| def _get_inbound_nodes(layer): | |
| """Returns a list of [name, size, index] for all inbound nodes of the given layer.""" | |
| inbound_nodes = [] | |
| if layer.get("inbound_nodes") is not None: | |
| for maybe_inbound_node in layer.get("inbound_nodes", []): | |
| if not isinstance(maybe_inbound_node, dict): | |
| # Note that the inbound node parsing is not backward compatible with | |
| # Keras 2. If given a Keras 2 model, the input nodes will be missing | |
| # in the final graph. | |
| continue | |
| for inbound_node_args in maybe_inbound_node.get("args", []): | |
| # Sometimes this field is a list when there are multiple inbound nodes | |
| # for the given layer. | |
| if not isinstance(inbound_node_args, list): | |
| inbound_node_args = [inbound_node_args] | |
| for arg in inbound_node_args: | |
| history = arg.get("config", {}).get("keras_history", []) | |
| if len(history) < 3: | |
| continue | |
| inbound_nodes.append(history[:3]) | |
| return inbound_nodes | |
| def _update_dicts( | |
| name_scope, | |
| model_layer, | |
| input_to_in_layer, | |
| model_name_to_output, | |
| prev_node_name, | |
| ): | |
| """Updates input_to_in_layer, model_name_to_output, and prev_node_name | |
| based on the model_layer. | |
| Args: | |
| name_scope: a string representing a scope name, similar to that of tf.name_scope. | |
| model_layer: a dict representing a Keras model configuration. | |
| input_to_in_layer: a dict mapping Keras.layers.Input to inbound layer. | |
| model_name_to_output: a dict mapping Keras Model name to output layer of the model. | |
| prev_node_name: a string representing a previous, in sequential model layout, | |
| node name. | |
| Returns: | |
| A tuple of (input_to_in_layer, model_name_to_output, prev_node_name). | |
| input_to_in_layer: a dict mapping Keras.layers.Input to inbound layer. | |
| model_name_to_output: a dict mapping Keras Model name to output layer of the model. | |
| prev_node_name: a string representing a previous, in sequential model layout, | |
| node name. | |
| """ | |
| layer_config = model_layer.get("config") | |
| if not layer_config.get("layers"): | |
| raise ValueError("layer is not a model.") | |
| node_name = _scoped_name(name_scope, layer_config.get("name")) | |
| input_layers = layer_config.get("input_layers") | |
| output_layers = layer_config.get("output_layers") | |
| inbound_nodes = _get_inbound_nodes(model_layer) | |
| is_functional_model = bool(input_layers and output_layers) | |
| # In case of [1] and the parent model is functional, current layer | |
| # will have the 'inbound_nodes' property. | |
| is_parent_functional_model = bool(inbound_nodes) | |
| if is_parent_functional_model and is_functional_model: | |
| for input_layer, inbound_node in zip(input_layers, inbound_nodes): | |
| input_layer_name = _scoped_name(node_name, input_layer) | |
| inbound_node_name = _scoped_name(name_scope, inbound_node[0]) | |
| input_to_in_layer[input_layer_name] = inbound_node_name | |
| elif is_parent_functional_model and not is_functional_model: | |
| # Sequential model can take only one input. Make sure inbound to the | |
| # model is linked to the first layer in the Sequential model. | |
| prev_node_name = _scoped_name(name_scope, inbound_nodes[0][0]) | |
| elif ( | |
| not is_parent_functional_model | |
| and prev_node_name | |
| and is_functional_model | |
| ): | |
| assert len(input_layers) == 1, ( | |
| "Cannot have multi-input Functional model when parent model " | |
| "is not Functional. Number of input layers: %d" % len(input_layer) | |
| ) | |
| input_layer = input_layers[0] | |
| input_layer_name = _scoped_name(node_name, input_layer) | |
| input_to_in_layer[input_layer_name] = prev_node_name | |
| if is_functional_model and output_layers: | |
| layers = _norm_to_list_of_layers(output_layers) | |
| layer_names = [_scoped_name(node_name, layer[0]) for layer in layers] | |
| model_name_to_output[node_name] = layer_names | |
| else: | |
| last_layer = layer_config.get("layers")[-1] | |
| last_layer_name = last_layer.get("config").get("name") | |
| output_node = _scoped_name(node_name, last_layer_name) | |
| model_name_to_output[node_name] = [output_node] | |
| return (input_to_in_layer, model_name_to_output, prev_node_name) | |
| def keras_model_to_graph_def(keras_layer): | |
| """Returns a GraphDef representation of the Keras model in a dict form. | |
| Note that it only supports models that implemented to_json(). | |
| Args: | |
| keras_layer: A dict from Keras model.to_json(). | |
| Returns: | |
| A GraphDef representation of the layers in the model. | |
| """ | |
| input_to_layer = {} | |
| model_name_to_output = {} | |
| g = GraphDef() | |
| # Sequential model layers do not have a field "inbound_nodes" but | |
| # instead are defined implicitly via order of layers. | |
| prev_node_name = None | |
| for name_scope, layer in _walk_layers(keras_layer): | |
| if _is_model(layer): | |
| ( | |
| input_to_layer, | |
| model_name_to_output, | |
| prev_node_name, | |
| ) = _update_dicts( | |
| name_scope, | |
| layer, | |
| input_to_layer, | |
| model_name_to_output, | |
| prev_node_name, | |
| ) | |
| continue | |
| layer_config = layer.get("config") | |
| node_name = _scoped_name(name_scope, layer_config.get("name")) | |
| node_def = g.node.add() | |
| node_def.name = node_name | |
| if layer.get("class_name") is not None: | |
| keras_cls_name = layer.get("class_name").encode("ascii") | |
| node_def.attr["keras_class"].s = keras_cls_name | |
| dtype_or_policy = layer_config.get("dtype") | |
| dtype = None | |
| has_unsupported_value = False | |
| # If this is a dict, try and extract the dtype string from | |
| # `config.name`. Keras will export like this for non-input layers and | |
| # some other cases (e.g. tf/keras/mixed_precision/Policy, as described | |
| # in issue #5548). | |
| if isinstance(dtype_or_policy, dict) and "config" in dtype_or_policy: | |
| dtype = dtype_or_policy.get("config").get("name") | |
| elif dtype_or_policy is not None: | |
| dtype = dtype_or_policy | |
| if dtype is not None: | |
| try: | |
| tf_dtype = dtypes.as_dtype(dtype) | |
| node_def.attr["dtype"].type = tf_dtype.as_datatype_enum | |
| except TypeError: | |
| has_unsupported_value = True | |
| elif dtype_or_policy is not None: | |
| has_unsupported_value = True | |
| if has_unsupported_value: | |
| # There's at least one known case when this happens, which is when | |
| # mixed precision dtype policies are used, as described in issue | |
| # #5548. (See https://keras.io/api/mixed_precision/). | |
| # There might be a better way to handle this, but here we are. | |
| logger.warning( | |
| "Unsupported dtype value in graph model config (json):\n%s", | |
| dtype_or_policy, | |
| ) | |
| if layer.get("inbound_nodes") is not None: | |
| for name, size, index in _get_inbound_nodes(layer): | |
| inbound_name = _scoped_name(name_scope, name) | |
| # An input to a layer can be output from a model. In that case, the name | |
| # of inbound_nodes to a layer is a name of a model. Remap the name of the | |
| # model to output layer of the model. Also, since there can be multiple | |
| # outputs in a model, make sure we pick the right output_layer from the model. | |
| inbound_node_names = model_name_to_output.get( | |
| inbound_name, [inbound_name] | |
| ) | |
| # There can be multiple inbound_nodes that reference the | |
| # same upstream layer. This causes issues when looking for | |
| # a particular index in that layer, since the indices | |
| # captured in `inbound_nodes` doesn't necessarily match the | |
| # number of entries in the `inbound_node_names` list. To | |
| # avoid IndexErrors, we just use the last element in the | |
| # `inbound_node_names` in this situation. | |
| # Note that this is a quick hack to avoid IndexErrors in | |
| # this situation, and might not be an appropriate solution | |
| # to this problem in general. | |
| input_name = ( | |
| inbound_node_names[index] | |
| if index < len(inbound_node_names) | |
| else inbound_node_names[-1] | |
| ) | |
| node_def.input.append(input_name) | |
| elif prev_node_name is not None: | |
| node_def.input.append(prev_node_name) | |
| if node_name in input_to_layer: | |
| node_def.input.append(input_to_layer.get(node_name)) | |
| prev_node_name = node_def.name | |
| return g | |