mct-nightly 1.1.0.7012022.post2611__py3-none-any.whl → 1.1.0.07122021-002414__py3-none-any.whl

model_compression_toolkit/common/graph/base_graph.py

@@ -21,12 +21,13 @@ import networkx as nx
 import numpy as np
 
 from networkx.algorithms.dag import topological_sort
+from model_compression_toolkit import common
 from model_compression_toolkit.common.graph.edge import EDGE_SINK_INDEX, EDGE_SOURCE_INDEX
 from model_compression_toolkit.common.graph.edge import Edge, convert_to_edge
 from model_compression_toolkit.common.graph.graph_searches import GraphSearches
-from model_compression_toolkit.common.graph.base_node import BaseNode
-from model_compression_toolkit.common.collectors.statistics_collector import BaseStatsCollector
-from model_compression_toolkit.common.collectors.statistics_collector import scale_statistics, shift_statistics
+from model_compression_toolkit.common.graph.node import Node
+from model_compression_toolkit.common.statistics_collector import BaseStatsContainer
+from model_compression_toolkit.common.statistics_collector import scale_statistics, shift_statistics
 from model_compression_toolkit.common.user_info import UserInformation
 from model_compression_toolkit.common.logger import Logger
 
@@ -39,8 +40,8 @@ class Graph(nx.MultiDiGraph, GraphSearches):
     """
 
     def __init__(self,
-                 nodes: List[BaseNode],
-                 input_nodes: List[BaseNode],
+                 nodes: List[Node],
+                 input_nodes: List[Node],
                  output_nodes: List[OutTensor],
                  edge_list: List[Edge],
                  **attr):
@@ -79,7 +80,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
 
         return np.unique([n.op for n in self.nodes()])
 
-    def get_inputs(self) -> List[BaseNode]:
+    def get_inputs(self) -> List[Node]:
         """
         Returns: List containing the model input nodes.
         """
@@ -94,7 +95,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         return self.output_nodes
 
     def set_inputs(self,
-                   input_nodes: List[BaseNode]):
+                   input_nodes: List[Node]):
         """
         Set the graph inputs dictionary.
         Args:
@@ -114,8 +115,8 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         self.output_nodes = output_nodes
 
     def set_out_stats_collector_to_node(self,
-                                        n: BaseNode,
-                                        stats_collector: BaseStatsCollector):
+                                        n: Node,
+                                        stats_collector: BaseStatsContainer):
         """
         Set an output statistics collector of a node in the graph, and set this statistics collector as an input
         statistics collector of nodes next to this given node.
@@ -163,7 +164,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
                     self.node_to_in_stats_collector.update({oe.sink_node: stats_collector})
 
     def get_out_stats_collector(self,
-                                n: BaseNode) -> BaseStatsCollector:
+                                n: Node) -> BaseStatsContainer:
         """
         Get the output statistics collector of a node containing output statistics of the node.
         Args:
@@ -172,10 +173,11 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         Returns:
             Tensor containing output statistics of the node.
         """
+
         return self.node_to_out_stats_collector.get(n)
 
     def get_in_stats_collector(self,
-                               n: BaseNode) -> BaseStatsCollector:
+                               n: Node) -> BaseStatsContainer:
         """
         Get the input statistics collector of a node containing input statistics of the node.
         Args:
@@ -191,7 +193,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         return sc
 
     def scale_stats_collector(self,
-                              node: BaseNode,
+                              node: Node,
                               scale_factor: np.ndarray):
         """
         Scale the output statistics of a node in the graph by a given scaling factor.
@@ -209,7 +211,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         self.set_out_stats_collector_to_node(node, scaled_sc)
 
     def shift_stats_collector(self,
-                              node: BaseNode,
+                              node: Node,
                               shift_value: np.ndarray):
         """
         Shift the output statistics of a node in the graph by a given value.
@@ -227,7 +229,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         self.set_out_stats_collector_to_node(node, shifted_sc)
 
     def find_node_by_name(self,
-                          name: str) -> List[BaseNode]:
+                          name: str) -> List[Node]:
         """
         Find and return a list of nodes by a name.
 
@@ -241,7 +243,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         return [n for n in self.nodes if n.name == name]
 
     def get_next_nodes(self,
-                       node_obj: BaseNode) -> List[BaseNode]:
+                       node_obj: Node) -> List[Node]:
         """
         Get next nodes (in a topological order) of a node.
 
@@ -256,7 +258,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         return [edges_list.sink_node for edges_list in self.out_edges(node_obj)]
 
     def get_prev_nodes(self,
-                       node_obj: BaseNode) -> List[BaseNode]:
+                       node_obj: Node) -> List[Node]:
         """
         Get previous nodes (in a topological order) of a node.
 
@@ -271,8 +273,8 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         return [edges_list.source_node for edges_list in self.incoming_edges(node_obj)]
 
     def reconnect_out_edges(self,
-                            current_node: BaseNode,
-                            new_node: BaseNode):
+                            current_node: Node,
+                            new_node: Node):
         """
         Connect all outgoing edges of a node to be outgoing edges of a different node
         (useful when replacing a node during substitutions).
@@ -287,8 +289,8 @@ class Graph(nx.MultiDiGraph, GraphSearches):
             self.remove_edge(current_node, oe.sink_node)
 
     def reconnect_in_edges(self,
-                           current_node: BaseNode,
-                           new_node: BaseNode):
+                           current_node: Node,
+                           new_node: Node):
         """
         Connect all incoming edges of a node to be incoming edges of a different node
         (useful when replacing a node during substitutions).
@@ -303,8 +305,8 @@ class Graph(nx.MultiDiGraph, GraphSearches):
             self.remove_edge(ie.source_node, current_node)
 
     def replace_output_node(self,
-                            current_node: BaseNode,
-                            new_node: BaseNode):
+                            current_node: Node,
+                            new_node: Node):
         """
         If a node is being substituted with another node and it is an output node, the graph's outputs
         should be updated as well. This function takes care of it by going over the graph's outputs, and
@@ -325,8 +327,8 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         self.set_outputs(new_graph_outputs)
 
     def remove_node(self,
-                    node_to_remove: BaseNode,
-                    new_graph_inputs: List[BaseNode] = None,
+                    node_to_remove: Node,
+                    new_graph_inputs: List[Node] = None,
                     new_graph_outputs: List[OutTensor] = None):
         """
         Remove a node from the graph. A new inputs/outputs lists can be passed in case the node is currently an
@@ -365,7 +367,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         super().remove_node(node_to_remove)
 
     def incoming_edges(self,
-                       n: BaseNode,
+                       n: Node,
                        sort_by_attr: str = None) -> List[Edge]:
         """
         Get a list of incoming edges of a node. If sort_by_attr is passed, the returned list
@@ -386,7 +388,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         return input_edges
 
     def out_edges(self,
-                  n: BaseNode,
+                  n: Node,
                   sort_by_attr: str = None) -> List[Edge]:
         """
         Get a list of outgoing edges of a node. If sort_by_attr is passed, the returned list
@@ -416,8 +418,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
             memory += n.get_memory_bytes()
         return memory
 
-    def get_configurable_sorted_nodes_names(self,
-                                            include_reused_nodes: bool = False) -> List[str]:
+    def get_configurable_sorted_nodes_names(self, include_reused_nodes: bool = False) -> List[str]:
         """
         Get a list of nodes' names that can be configured (namely, has one or
         more weight qc candidate). The names are sorted according to the topological
@@ -430,15 +431,14 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         more weight qc candidate) sorted topology.
 
         """
-        sorted_names = [n.name for n in self.get_configurable_sorted_nodes(include_reused_nodes=include_reused_nodes)]
+        sorted_names = [n.name for n in self.get_configurable_sorted_nodes(include_reused_nodes)]
         return sorted_names
 
-    def get_configurable_sorted_nodes(self,
-                                      include_reused_nodes: bool = False) -> List[BaseNode]:
+    def get_configurable_sorted_nodes(self, include_reused_nodes: bool = False) -> List[Node]:
         """
         Get a list of nodes that can be configured (namely, has one or
-        more weight qc candidate and their weights should be quantized).
-        The nodes are sorted according to the topological order of the graph.
+        more weight qc candidate). The nodes are sorted according to the topological
+        order of the graph.
 
         Args:
             include_reused_nodes: Whether or not to include reused nodes (False by default).
@@ -450,7 +450,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         sorted_configurable_nodes = []
         sorted_nodes = list(topological_sort(self))
         for n in sorted_nodes:
-            if n.is_weights_quantization_enabled():
+            if n.candidates_weights_quantization_cfg is not None:
                 if not n.reuse or include_reused_nodes:
                     if len(n.candidates_weights_quantization_cfg) >= 1:
                         sorted_configurable_nodes.append(n)

model_compression_toolkit/common/graph/edge.py

@@ -16,7 +16,7 @@
 
 from typing import Any, Dict
 
-from model_compression_toolkit.common.graph.base_node import BaseNode
+from model_compression_toolkit.common.graph.node import Node
 
 # Edge attributes:
 EDGE_SOURCE_INDEX = 'source_index'
@@ -29,8 +29,8 @@ class Edge(object):
     """
 
     def __init__(self,
-                 source_node: BaseNode,
-                 sink_node: BaseNode,
+                 source_node: Node,
+                 sink_node: Node,
                  source_index: int,
                  sink_index: int):
         """

model_compression_toolkit/common/graph/graph_matchers.py

@@ -18,7 +18,7 @@ from typing import Any, List
 
 from tensorflow.keras.layers import Layer
 
-from model_compression_toolkit.common.graph.base_node import BaseNode
+from model_compression_toolkit.common.graph.node import Node
 from model_compression_toolkit.common.matchers import node_matcher, walk_matcher, edge_matcher
 
 
@@ -92,7 +92,7 @@ class EdgeMatcher(edge_matcher.BaseEdgeMatcher):
     class EdgeMatcher to check if an edge matches an edge that EdgeMatcher contains.
     """
 
-    def __init__(self, source_matcher: BaseNode, target_matcher: BaseNode):
+    def __init__(self, source_matcher: Node, target_matcher: Node):
         """
         Init an EdgeMatcher object.
 
@@ -125,7 +125,7 @@ class WalkMatcher(walk_matcher.WalkMatcherList):
     Class WalkMatcher to check if a list of nodes matches another list of nodes.
     """
 
-    def __init__(self, matcher_list: List[BaseNode]):
+    def __init__(self, matcher_list: List[Node]):
         """
         Init a WalkMatcher object.
 

model_compression_toolkit/common/graph/graph_searches.py

@@ -16,7 +16,7 @@
 from abc import ABC
 from typing import List, Any
 
-from model_compression_toolkit.common.graph.base_node import BaseNode
+from model_compression_toolkit.common.graph.node import Node
 from model_compression_toolkit.common.matchers import node_matcher, base_graph_filter, edge_matcher
 from model_compression_toolkit.common.matchers.walk_matcher import WalkMatcherList
 
@@ -59,7 +59,7 @@ class GraphSearches(base_graph_filter.BaseGraphFilter, ABC):
 
         return edge_list
 
-    def _walk_filter(self, walk_matcher: WalkMatcherList) -> List[BaseNode]:
+    def _walk_filter(self, walk_matcher: WalkMatcherList) -> List[Node]:
         """
         Search for a list of nodes which match the list in walk_matcher.
         If one the nodes in the list (that was found in the graph) has more than one output,
@@ -72,8 +72,8 @@ class GraphSearches(base_graph_filter.BaseGraphFilter, ABC):
             A list of nodes which match the list in walk_matcher.
         """
 
-        def walk_match(node: BaseNode,
-                       node_list: List[BaseNode],
+        def walk_match(node: Node,
+                       node_list: List[Node],
                        index: int,
                        node_matcher_list: list) -> Any:
             """

model_compression_toolkit/common/graph/graph_vis.py

@@ -0,0 +1,116 @@
+# Copyright 2021 Sony Semiconductors Israel, Inc. 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.
+# ==============================================================================
+
+
+import json
+import os
+from typing import Any, Dict, Tuple
+
+from model_compression_toolkit import common
+from model_compression_toolkit.common.graph.node import Node
+
+
+def check_str(in_sting: str) -> bool:
+    """
+    Checks if an open bracket is in a string.
+
+    Args:
+        in_sting: String to check.
+
+    Returns:
+        Whether an open bracket is in the string or not.
+    """
+
+    return '(' not in in_sting and '[' not in in_sting and '{' not in in_sting
+
+
+def node_dict(n: Node) -> Dict[str, Any]:
+    """
+    Get a dictionary with a node's attributes for displaying.
+
+    Args:
+        n: Node to get its attributes to display.
+
+    Returns:
+        A dictionary with params of the node to display when visualizing the graph.
+    """
+
+    framework_attr = {k: str(v) for k, v in n.framework_attr.items() if check_str(str(v))}
+    framework_attr.update({k: str(v) for k, v in n.quantization_attr.items() if check_str(str(v))})
+    framework_attr.update({'op': n.layer_class.__name__})
+
+    if n.quantization_cfg is not None:
+        for k, v in n.quantization_cfg.activation_quantization_params.items():
+            framework_attr.update({k: str(v)})
+        framework_attr.update({'activation_is_signed': str(n.quantization_cfg.activation_is_signed)})
+
+    return {"id": n.name,
+            "group": n.layer_class.__name__,
+            "label": n.layer_class.__name__,
+            "title": "",
+            "properties": framework_attr}
+
+
+def edge_dict(i: int,
+              edge: Tuple[Node, Node],
+              graph):
+    """
+    Create a dictionary of attributes to visualize an edge in the graph.
+
+    Args:
+        i: Edge's ID.
+        edge: Tuple of two nodes (source and destination).
+        graph: Graph the edge is in.
+
+    Returns:
+        Dictionary of attributes to visualize an edge in a graph.
+    """
+
+    return {"id": str(i),
+            "from": edge[0].name,
+            "to": edge[1].name,
+            "label": "",
+            "title": "",
+            "properties": graph.get_edge_data(edge[0], edge[1])}
+
+
+def write_vis_graph(folder_path: str,
+                    file_name: str,
+                    graph):
+    """
+    Create and save a json file containing data about the graph to visualize (such as nodes and edges).
+    folder_path and file_name determine where the json file will be saved.
+
+    Args:
+        folder_path: Dir path to save the json file.
+        file_name: File name of the json file.
+        graph: Graph to visualize.
+
+    """
+
+    os.makedirs(folder_path, exist_ok=True)
+    file_path = os.path.join(folder_path, file_name + '.json')
+
+    nodes_list = [node_dict(n) for n in graph.nodes()]
+    edges_list = [edge_dict(i, e, graph) for i, e in enumerate(graph.edges())]
+    graph_dict = {'nodes': nodes_list,
+                  'edges': edges_list,
+                  "options": {}
+                  }
+
+    with open(file_path, 'w') as json_file:
+        json.dump(graph_dict, json_file)
+
+    common.Logger.info(f"Writing Vis Graph to:{file_path}")

model_compression_toolkit/common/graph/{base_node.py → node.py}

@@ -21,7 +21,7 @@ import numpy as np
 from model_compression_toolkit.common.constants import WEIGHTS_NBITS_ATTRIBUTE, CORRECTED_BIAS_ATTRIBUTE
 
 
-class BaseNode:
+class Node:
     """
     Class to represent a node in a graph that represents the model.
     """
@@ -35,6 +35,7 @@ class BaseNode:
                  layer_class: type,
                  reuse: bool = False,
                  reuse_group: str = None,
+                 op_call_args: Dict[str, Any] = {},
                  quantization_attr: Dict[str, Any] = None):
         """
         Init a Node object.
@@ -48,8 +49,11 @@ class BaseNode:
             layer_class: Class path of the layer this node represents.
             reuse: Whether this node was duplicated and represents a reused layer.
             reuse_group: Name of group of nodes from the same reused layer.
+            op_call_args: Arguments dictionary with values to pass when calling the layer.
             quantization_attr: Attributes the node holds regarding how it should be quantized.
         """
+
+
         self.name = name
         self.framework_attr = framework_attr
         self.quantization_attr = quantization_attr if quantization_attr is not None else dict()
@@ -62,25 +66,37 @@ class BaseNode:
         self.activation_quantization_cfg = None
         self.final_weights_quantization_cfg = None
         self.candidates_weights_quantization_cfg = None
-        self.prior_info = None
+        self.output_quantization = True
+        self.op_call_args = op_call_args
+
+    def no_quantization(self) -> bool:
+        """
+
+        Returns: Whether NodeQuantizationConfig does not have activation params.
+
+        """
+        return self.activation_quantization_cfg is None or \
+               (not self.activation_quantization_cfg.has_activation_quantization_params())
 
-    def is_activation_quantization_enabled(self) -> bool:
+    def weight_quantization(self) -> bool:
         """
 
-        Returns: Whether node activation quantization is enabled or not.
+        Returns: Whether node weights should be quantized
 
         """
-        return self.activation_quantization_cfg.enable_activation_quantization
+        return self.final_weights_quantization_cfg is not None and \
+               self.final_weights_quantization_cfg.has_weights_quantization_params() and \
+               self.final_weights_quantization_cfg.enable_weights_quantization
 
-    def is_weights_quantization_enabled(self) -> bool:
+    def activation_quantization(self) -> bool:
         """
 
-        Returns: Whether node weights quantization is enabled or not.
+        Returns: Whether node activation should be quantized
 
         """
-        for qc in self.candidates_weights_quantization_cfg:
-            assert self.candidates_weights_quantization_cfg[0].enable_weights_quantization == qc.enable_weights_quantization
-        return self.candidates_weights_quantization_cfg[0].enable_weights_quantization
+        return self.activation_quantization_cfg is not None and \
+               self.activation_quantization_cfg.has_activation_quantization_params() and \
+               self.activation_quantization_cfg.enable_activation_quantization
 
     def __repr__(self):
         """
@@ -172,7 +188,7 @@ class BaseNode:
         """
         shared_attributes = [CORRECTED_BIAS_ATTRIBUTE, WEIGHTS_NBITS_ATTRIBUTE]
         attr = dict()
-        if self.is_weights_quantization_enabled():
+        if self.candidates_weights_quantization_cfg is not None:
             attr = copy.deepcopy(self.candidates_weights_quantization_cfg[0].__dict__)
             for shared_attr in shared_attributes:
                 if shared_attr in attr:

model_compression_toolkit/common/mixed_precision/bit_width_setter.py

@@ -18,13 +18,14 @@ import copy
 from typing import Any, List
 
 from model_compression_toolkit.common.quantization.quantization_config import QuantizationConfig
-from model_compression_toolkit.common import Graph, BaseNode
+from model_compression_toolkit.common import Graph, Node
 from model_compression_toolkit.common.framework_info import FrameworkInfo
 from model_compression_toolkit.common.logger import Logger
 from model_compression_toolkit.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfig
 
 
+
 def set_bit_widths(quant_config: QuantizationConfig,
                    graph_to_set_bit_widths: Graph,
                    fw_info: FrameworkInfo = None,
@@ -36,10 +37,9 @@ def set_bit_widths(quant_config: QuantizationConfig,
     Args:
         quant_config: MixedPrecisionQuantizationConfig the graph was computed according to.
         graph_to_set_bit_widths: A prepared for quantization graph to set its bit widths.
-        fw_info: Information needed for quantization about the specific framework (e.g., kernel
-        channels indices, groups of layers by how they should be quantized, etc.)
-        bit_widths_config: MP configuration (a list of indices: one for each node's candidate
-        quantization configuration).
+        fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
+        groups of layers by how they should be quantized, etc.)
+        bit_widths_config: MP configuration (a list of indices: one for each node's candidate quantization configuration).
 
     """
     graph = copy.deepcopy(graph_to_set_bit_widths)
@@ -63,8 +63,7 @@ def set_bit_widths(quant_config: QuantizationConfig,
             # Get a list of nodes' names we need to finalize (that they have at least one weight qc candidate).
             sorted_nodes_names = graph.get_configurable_sorted_nodes_names()
             for node in graph.nodes:  # set a specific node qc for each node final weights qc
-                # If it's reused, take the configuration that the base node has
-                node_name = node.name if not node.reuse else '_'.join(node.name.split('_')[:-2])
+                node_name = node.name if not node.reuse else '_'.join(node.name.split('_')[:-2]) # if it's reused, take the configuration that the base node has
                 if node_name in sorted_nodes_names:  # only configurable nodes are in this list
                     node_index_in_graph = sorted_nodes_names.index(node_name)
                     _set_node_qc(bit_widths_config,
@@ -83,7 +82,7 @@ def set_bit_widths(quant_config: QuantizationConfig,
     return graph
 
 
-def _get_node_qc_by_bit_widths(node: BaseNode,
+def _get_node_qc_by_bit_widths(node: Node,
                                bit_width_cfg: List[int],
                                node_index_in_graph: int) -> Any:
     """
@@ -100,7 +99,7 @@ def _get_node_qc_by_bit_widths(node: BaseNode,
         Node quantization configuration if it was found, or None otherwise.
     """
 
-    if node.is_weights_quantization_enabled():
+    if node.candidates_weights_quantization_cfg is not None:
         bit_index_in_cfg = bit_width_cfg[node_index_in_graph]
         qc = node.candidates_weights_quantization_cfg[bit_index_in_cfg]
         return qc
@@ -109,7 +108,7 @@ def _get_node_qc_by_bit_widths(node: BaseNode,
 
 def _set_node_qc(bit_width_cfg: List[int],
                  fw_info: FrameworkInfo,
-                 node: BaseNode,
+                 node: Node,
                  node_index_in_graph: int):
     """
     Get the node's quantization configuration that

model_compression_toolkit/common/mixed_precision/mixed_precision_search_manager.py

@@ -111,7 +111,7 @@ class MixedPrecisionSearchManager(object):
                 if n.name in mp_nodes:
                     node_idx = mp_nodes.index(n.name)
                     node_nbits = n.candidates_weights_quantization_cfg[mp_model_config[node_idx]].weights_n_bits
-                elif n.is_weights_quantization_enabled():
+                elif n.candidates_weights_quantization_cfg is not None:
                     # The only valid way to get here is if the node is reused (which means that we're not looking
                     # for its configuration), and we ignore it when computing the KPI (as the base node will acount
                     # for it).

model_compression_toolkit/common/model_collector.py

@@ -17,7 +17,7 @@
 import numpy as np
 from typing import List
 
-from model_compression_toolkit import FrameworkInfo, common
+from model_compression_toolkit import FrameworkInfo
 from model_compression_toolkit.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.common.graph.base_graph import Graph
 from model_compression_toolkit.common.logger import Logger
@@ -53,17 +53,15 @@ class ModelCollector(object):
 
         for n in self.graph.nodes():
             out_stats_container = self.graph.get_out_stats_collector(n)
+
             if isinstance(out_stats_container, list):  # If layer has multiple outputs
-                # Append nodes to output and track their statistics only if
-                # they actually collect statistics.
-                if len([x for x in out_stats_container if not isinstance(x, common.NoStatsCollector)]) > 0:
-                    mark2fetch = True
-                    for sc in out_stats_container:
-                        # Output only if statistics should be gathered
-                        if sc.require_collection() and mark2fetch:
-                            mark2fetch = False
-                            node2fetch.append(n)  # Append node several times (as number of outputs it has)
-                    stats_containers_list.append(out_stats_container)
+                mark2fetch = True
+                for sc in out_stats_container:
+                    # Output only if statistics should be gathered
+                    if sc.require_collection() and mark2fetch:
+                        mark2fetch = False
+                        node2fetch.append(n)  # Append node several times (as number of outputs it has)
+                stats_containers_list.append(out_stats_container)
 
             else:  # A single output
                 if out_stats_container.require_collection():
@@ -75,8 +73,8 @@ class ModelCollector(object):
         # Build a float model and output all layers' outputs
         # (that should be collected) as the model's outputs
         self.model, _ = self.fw_impl.model_builder(self.graph,
-                                                   mode=ModelBuilderMode.FLOAT,
-                                                   append2output=node2fetch,
+                                                      mode=ModelBuilderMode.FLOAT,
+                                                      append2output=node2fetch,
                                                    fw_info=self.fw_info)
 
     def infer(self, inputs_list: List[np.ndarray]):
@@ -102,4 +100,4 @@ class ModelCollector(object):
                 for tdi, sci in zip(td, sc):
                     sci.update_statistics(self.fw_impl.to_numpy(tdi))
             else:
-                sc.update_statistics(self.fw_impl.to_numpy(td))
+                sc.update_statistics(self.fw_impl.to_numpy(td))