diff --git a/nncf/openvino/graph/model_utils.py b/nncf/openvino/graph/model_utils.py index 22de60ed0e3..084ffc9a79b 100644 --- a/nncf/openvino/graph/model_utils.py +++ b/nncf/openvino/graph/model_utils.py @@ -59,4 +59,4 @@ def get_input_nodes(nncf_graph: NNCFGraph) -> List[NNCFNode]: :param nncf_graph: NNCFGraph to work with. :return: Target NNCFGraph input nodes. """ - return list(set(nncf_graph.get_input_nodes()).upadte(nncf_graph.get_nodes_by_metatypes([OVReadValueMetatype]))) + return nncf_graph.get_input_nodes() + nncf_graph.get_nodes_by_metatypes([OVReadValueMetatype]) diff --git a/nncf/openvino/graph/nncf_graph_builder.py b/nncf/openvino/graph/nncf_graph_builder.py index ea8cb9b0e37..80e109a7846 100644 --- a/nncf/openvino/graph/nncf_graph_builder.py +++ b/nncf/openvino/graph/nncf_graph_builder.py @@ -199,8 +199,10 @@ def _set_non_weighted_layer_attributes(node: ov.Node, metatype: OVOpMetatype, nn """ if metatype == OVConcatMetatype: nncf_node = nncf_graph.get_node_by_name(node.get_friendly_name()) + axis = axis = node.get_attributes()["axis"] + num_inputs = len(node.inputs()) nncf_node.layer_attributes = OVLayerAttributes( - {}, MultipleInputLayerAttributes(axis=node.get_axis(), num_inputs=len(node.inputs())) + {}, MultipleInputLayerAttributes(axis=axis, num_inputs=num_inputs) ) def _set_weighted_layer_attributes( diff --git a/tests/openvino/native/test_layer_attributes.py b/tests/openvino/native/test_layer_attributes.py index 7dacec14159..8f786d30231 100644 --- a/tests/openvino/native/test_layer_attributes.py +++ b/tests/openvino/native/test_layer_attributes.py @@ -160,13 +160,21 @@ def get_lstm(node_name, input_shape): return [input_1], [output] -def get_concat_node(op_name, input_shape): +def get_concat_node(op_name, input_shape, axis): num_inputs = 3 inputs = [opset.parameter(input_shape, name=f"Input{i}") for i in range(num_inputs)] - output = opset.concat(inputs, axis=1, name=op_name) + output = opset.concat(inputs, axis=axis, name=op_name) return inputs, [output] +def get_concat_axis_default(op_name, input_shape): + return get_concat_node(op_name, input_shape, axis=0) + + +def get_concat_axis_custom(op_name, input_shape): + return get_concat_node(op_name, input_shape, axis=1) + + def get_one_layer_model(op_name: str, node_creator, input_shape): inputs, outputs = node_creator(op_name, input_shape) results = [opset.result(o, name=f"Result{i}") for i, o in enumerate(outputs)] @@ -446,7 +454,10 @@ class LayerAttributesTestCase: None, ), LayerAttributesTestCase( - get_concat_node, (1, 2, 3, 3), OVLayerAttributes({}, MultipleInputLayerAttributes(1, 3)), None + get_concat_axis_default, (1, 2, 3, 3), OVLayerAttributes({}, MultipleInputLayerAttributes(0, 3)), None + ), + LayerAttributesTestCase( + get_concat_axis_custom, (1, 2, 3, 3), OVLayerAttributes({}, MultipleInputLayerAttributes(1, 3)), None ), ]