leap-model-parser 0.1.223__tar.gz → 0.1.225.dev1__tar.gz

{leap_model_parser-0.1.223 → leap_model_parser-0.1.225.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.223
+Version: 0.1.225.dev1
 Summary: 
 Home-page: https://github.com/tensorleap/leap-model-parser
 License: MIT

{leap_model_parser-0.1.223 → leap_model_parser-0.1.225.dev1}/leap_model_parser/contract/graph.py

@@ -47,6 +47,7 @@ class Node:
     outputs: Dict[str, ConnectionOutput] = field(default_factory=dict)
     pruning_plan_id: Optional[str] = None
     wrapper: Optional[WrapperData] = None
+    shape: Optional[List[Optional[int]]] = None
 
     def __key(self):
         return (self.id, self.name)

{leap_model_parser-0.1.223 → leap_model_parser-0.1.225.dev1}/leap_model_parser/keras_json_model_import.py

@@ -5,7 +5,6 @@ from typing import Set, Dict, Any, List, Type, Optional, Tuple
 
 import tensorflow as tf  # type: ignore
 from keras.engine.keras_tensor import KerasTensor  # type: ignore
-from keras.engine.node import Node  # type: ignore
 from keras.engine.node import Node as keras_node # type: ignore
 from keras.layers import Layer # type: ignore
 from keras.layers.convolutional.base_conv import Conv # type: ignore
@@ -221,7 +220,7 @@ class KerasJsonModelImport:
 
         self.layer_data_adjustments(data, layer_metadata)
         node = Node(id=str(self.id), name=layer_metadata.get(
-            "class_name", layer["class_name"]), data=data, position=[0, 0])
+            "class_name", layer["class_name"]), data=data, position=[0, 0], shape=layer['output_shape'])
         if 'wrapper' in layer:
             node.wrapper = layer['wrapper']
         self.nodes_cache[node_key] = node

{leap_model_parser-0.1.223 → leap_model_parser-0.1.225.dev1}/leap_model_parser/leap_graph_editor.py

@@ -49,6 +49,7 @@ class LeapGraphEditor:
                 self._add_connection_to_node(new_metric_node_id, input_name, input_node_id)
         elif node_connection.node.type in (NodeMappingType.Loss, NodeMappingType.CustomLoss):
             new_loss_node_id = self._add_loss_node(node_connection.node.name,
+                                                   node_connection.node.user_unique_name,
                                                    node_connection.node.type == NodeMappingType.CustomLoss,
                                                    node_connection.node.arg_names)
             for input_name, node in node_connection.node_inputs.items():
@@ -159,7 +160,7 @@ class LeapGraphEditor:
         self.model_graph[new_node_id] = metric_node
         return new_node_id
 
-    def _add_loss_node(self, loss_name: str, is_custom_loss: bool, arg_names: Optional[List[str]]=None) -> str:
+    def _add_loss_node(self, loss_name: str, user_unique_name:str, is_custom_loss: bool, arg_names: Optional[List[str]]=None) -> str:
         new_node_id = self._generate_new_node_id()
 
         loss_type = 'CustomLoss' if is_custom_loss else 'Loss'
@@ -169,7 +170,7 @@ class LeapGraphEditor:
             new_node_id,
             loss_node_name,
             position=[0, 0],
-            data={'type': loss_type, 'selected': loss_name, 'name': loss_name},
+            data={'type': loss_type, 'selected': loss_name, 'name': loss_name, 'user_unique_name': user_unique_name},
             inputs={},
             outputs={
                 f'{new_node_id}-loss': ConnectionOutput([])

{leap_model_parser-0.1.223 → leap_model_parser-0.1.225.dev1}/leap_model_parser/model_parser.py

@@ -54,6 +54,10 @@ class ModelParser:
             ImportModelTypeEnum.PB_TF2.value: self.convert_pb_model,
         }
 
+    def _add_output_node_shape_to_model_schema(self, model_schema: Dict, keras_model: Model):
+        for i, layer in enumerate(keras_model.layers):
+            model_schema['config']['layers'][i]['output_shape'] = list(layer.output_shape)
+
     def get_keras_model_and_model_graph(
             self, model_path: Path, model_type: ImportModelTypeEnum) -> Tuple[Dict[str, Node], List[InputInfo], Optional[Model], Optional[str]]:
         model_to_keras_converter: Optional[Callable[[str], Tuple[Dict[str, Node], Model, Optional[str]]]] = \
@@ -64,6 +68,9 @@ class ModelParser:
 
         file_path = str(model_path)
         model_schema, keras_model_with_weights, error_info = model_to_keras_converter(file_path)
+
+        self._add_output_node_shape_to_model_schema(model_schema, keras_model_with_weights)
+
         model_generator = KerasJsonModelImport(self.custom_layers)
 
         keras_model = keras_model_with_weights

{leap_model_parser-0.1.223 → leap_model_parser-0.1.225.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leap-model-parser"
-version = "0.1.223"
+version = "0.1.225.dev1"
 description = ""
 authors = ["idan <idan.yogev@tensorleap.ai>"]
 license = "MIT"