U Vc(@sPdZddlmZddlmZddZddZdd Zd d Zd d Z ddZ dS)aJUtilities 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. )GraphDef)dtypesccstd|fV|ddrp|dd}|ddD]4}t|D]&\}}|r^d||fn|}||fVqFq:dS)aYWalks 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. configlayersname%s/%sN)get _walk_layers) keras_layer name_scopelayerZsub_name_scopeZsublayerrH/tmp/pip-unpacked-wheel-g8kmtpbc/tensorboard/plugins/graph/keras_util.pyr 0s r cCs|rd||fS|S)a/Returns scoped name for a node as a string in the form '/'. 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. rr)r node_namerrr _scoped_nameHs  rcCs|dddk S)zReturns True if layer is a model. Args: layer: a dict representing a Keras model configuration. Returns: bool: True if layer is a model. rrN)r )r rrr _is_modelXs rcCst|dtfr|S|gS)aNormalizes 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]. r) isinstancelist)Z maybe_layersrrr_norm_to_list_of_layersdsrcsj|d}|dstdt||d|d}|d}|d}t|oR|} t|} | r| rt||D](\} } t| } t|| d}||| <qpnb| r| st||ddd}nB| s|r| rt|d kstd t| |d} t| } ||| <| r.|r.t|}fd d |D}||<n2|dd }|dd}t|}|g|<|||fS)aUpdates 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. rrzlayer is not a model.r input_layers output_layers inbound_nodesrzhCannot have multi-input Functional model when parent model is not Functional. Number of input layers: %dcsg|]}t|dqS)r)r).0r rrr sz!_update_dicts..)r ValueErrorrboolziplenAssertionErrorr)r Z model_layerZinput_to_in_layermodel_name_to_outputprev_node_name layer_configrrrZis_functional_modelZis_parent_functional_modelZ input_layerZ inbound_nodeZinput_layer_nameZinbound_node_namerZ layer_namesZ last_layerZlast_layer_nameZ output_noderrr _update_dictsxsN             r&cCsti}i}t}d}t|D]R\}}t|rDt|||||\}}}q|d}t||d}|j} || _|ddk r|d d} | | j d_ |d} | dk rt | t st|d} | j| j d_|ddk r6|dD]F} t| }|D]4\}}}}t||}|||g}| j||qqn|dk rL| j|||krh| j||| j}q|S) aReturns 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. Nrr class_nameasciiZ keras_classZdtyper)rr rr&r rnodeaddrencodeattrsrdictrZas_dtypeZas_datatype_enumtyperinputappend)r Zinput_to_layerr#gr$r r r%rZnode_defZkeras_cls_nameZdtype_or_policyZtf_dtypeZmaybe_inbound_noderrsizeindex_Z inbound_nameZinbound_node_namesrrrkeras_model_to_graph_defs`          r6N) __doc__Z"tensorboard.compat.proto.graph_pb2rZ"tensorboard.compat.tensorflow_stubrr rrrr&r6rrrrs   I