mct-nightly 1.10.0.20231211.post417__py3-none-any.whl → 1.10.0.20231213.post410__py3-none-any.whl

{mct_nightly-1.10.0.20231211.post417.dist-info → mct_nightly-1.10.0.20231213.post410.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 1.10.0.20231211.post417
+Version: 1.10.0.20231213.post410
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct_nightly-1.10.0.20231211.post417.dist-info → mct_nightly-1.10.0.20231213.post410.dist-info}/RECORD

@@ -32,7 +32,7 @@ model_compression_toolkit/core/common/collectors/statistics_collector_generator.
 model_compression_toolkit/core/common/fusion/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/common/fusion/layer_fusing.py,sha256=tIsWFYc771o59uvq5fxAaBmOCnd_gd-_xMbQI9SupQA,5479
 model_compression_toolkit/core/common/graph/__init__.py,sha256=Xr-Lt_qXMdrCnnOaUS_OJP_3iTTGfPCLf8_vSrQgCs0,773
-model_compression_toolkit/core/common/graph/base_graph.py,sha256=-YBWWl3ZS7FJvZldGqT4SetlKI8j60f4sS0rYyFQpbI,30059
+model_compression_toolkit/core/common/graph/base_graph.py,sha256=UHhXCWXh4hK7cg2sZrNLiTYSYsvDc3yhImcdKbQ0VVs,30929
 model_compression_toolkit/core/common/graph/base_node.py,sha256=csIgi5ex7EquQsF34w5waRIHzbg7XitvIqQgCC29azs,21118
 model_compression_toolkit/core/common/graph/edge.py,sha256=K6Wc2hBcIqig5PbbLhbjtTgYtkyZEohfgj4Wn_J5yEA,3733
 model_compression_toolkit/core/common/graph/functional_node.py,sha256=0TpYNa2ODZ0M9lQ2z_GsStqAbrg1Muwdni74LjphAh0,2922
@@ -197,7 +197,7 @@ model_compression_toolkit/core/pytorch/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKW
 model_compression_toolkit/core/pytorch/constants.py,sha256=Kt_GDwe3yX9oMS1DI2eXYuUT25_lpjeCkxpstsAiXCI,2472
 model_compression_toolkit/core/pytorch/default_framework_info.py,sha256=qee8TFcDro2lfyXe_fujjX2OlxELTyKSsLlZ7QkzeXU,4200
 model_compression_toolkit/core/pytorch/kpi_data_facade.py,sha256=J0IDOtFMVFSFyBXDzNGbwJfHu89iRBJFdid1_wFB-xQ,8482
-model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=oTAd6_XYtyvTX2fRXx0BzajvgqbYreXGKD7ij8iL2SY,26482
+model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=w4yLHjJmOxfkSgApEx9rWAEbv9vkLnZik5JQvaX55FM,26654
 model_compression_toolkit/core/pytorch/pytorch_node_prior_info.py,sha256=n_B4a6FMwM9D2w8kzy3oenBWZgXNZuIZgTJC6JEuTy0,3250
 model_compression_toolkit/core/pytorch/utils.py,sha256=rBQMAbWluyIMjVfeghzq6FZv3sR_khszSRpbWvwussw,2959
 model_compression_toolkit/core/pytorch/back2framework/__init__.py,sha256=H_WixgN0elVWf3exgGYsi58imPoYDj5eYPeh6x4yfug,813
@@ -216,6 +216,7 @@ model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/batchno
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/batchnorm_reconstruction.py,sha256=B7aC2TZNrQJ2oQVGBFhKAVqdUU5lYVJSMmwKhjxOHWk,2822
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/batchnorm_refusing.py,sha256=JDWOaNwYrZG0zTwd3HwoZUM3tKu7zPbzLOrqNQsu8xA,2162
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/const_holder_conv.py,sha256=4mnowFmfDQjKlhHqsNto1iL4WbHyh4cM3Lf67Z-Cnzc,4804
+model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_batch_norm.py,sha256=bSexq2A7WxLLm13v67SgVbb4T1Y6nrKQDZfk4iSj_ec,3941
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/linear_collapsing.py,sha256=qCNT3L4mnZtIP75c8YwImvsTWdPIdsEvO4pc3SE4y6s,5797
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/multi_head_attention_decomposition.py,sha256=gGlVE1xUa1Pv2NbRx0I2y5Okg3kneBWSx9JwULuTWz0,38353
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/permute_call_method.py,sha256=EMCviyFyJFLEKuAUz3rZHLfB9MAU1kywSBL2XQNzLlg,1953
@@ -450,8 +451,8 @@ model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py,sha
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=MVwXNymmFRB2NXIBx4e2mdJ1RfoHxRPYRgjb1MQP5kY,1797
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=SbvRlIdE32PEBsINt1bhSqvrKL_zbM9V-aeSkOn-sw4,3083
-mct_nightly-1.10.0.20231211.post417.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-1.10.0.20231211.post417.dist-info/METADATA,sha256=krO3OedYSfl2Ck74IQ5o9-Bwi7174jQvnoy_WSEC2JE,16232
-mct_nightly-1.10.0.20231211.post417.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-mct_nightly-1.10.0.20231211.post417.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-1.10.0.20231211.post417.dist-info/RECORD,,
+mct_nightly-1.10.0.20231213.post410.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-1.10.0.20231213.post410.dist-info/METADATA,sha256=udqlXh19QUWj8Kl3urGNDSPuK2qjV7ZEY_8OFS-gtvw,16232
+mct_nightly-1.10.0.20231213.post410.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+mct_nightly-1.10.0.20231213.post410.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-1.10.0.20231213.post410.dist-info/RECORD,,

model_compression_toolkit/core/common/graph/base_graph.py

@@ -299,19 +299,24 @@ 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: BaseNode,
+                       sink_index_sorted: bool = False) -> List[BaseNode]:
         """
         Get previous nodes (in a topological order) of a node.
 
         Args:
             node_obj: Node to get its previous nodes.
+            sink_index_sorted: Whether to sort the returned list by the sink_index of the edges.
 
         Returns:
             List of input nodes objects.
 
         """
-
-        return [edges_list.source_node for edges_list in self.incoming_edges(node_obj)]
+        if sink_index_sorted:
+            sort_attr = 'sink_index'
+        else:
+            sort_attr = None
+        return [edges_list.source_node for edges_list in self.incoming_edges(node_obj, sort_by_attr=sort_attr)]
 
     def reconnect_out_edges(self,
                             current_node: BaseNode,
@@ -705,3 +710,19 @@ class Graph(nx.MultiDiGraph, GraphSearches):
 
         """
         return all([n.is_all_activation_candidates_equal() for n in self.nodes])
+
+    def replace_node(self, node_to_replace: BaseNode, new_node: BaseNode):
+        """
+        Replaces a node in the graph with a new node.
+
+        Args:
+            node_to_replace: The node to replace.
+            new_node: The new node to replace with.
+
+        """
+        self.add_node(new_node)
+        self.reconnect_out_edges(node_to_replace, new_node)
+        self.reconnect_in_edges(node_to_replace, new_node)
+        self.replace_output_node(node_to_replace, new_node)
+        self.replace_input_node(node_to_replace, new_node)
+        self.remove_node(node_to_replace)

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_batch_norm.py

@@ -0,0 +1,94 @@
+# Copyright 2023 Sony Semiconductor 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.
+# ==============================================================================
+from torch import nn
+import torch.nn.functional as F
+
+from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
+from model_compression_toolkit.core import common
+from model_compression_toolkit.core.common import BaseNode, Graph
+from model_compression_toolkit.core.pytorch.constants import *
+from model_compression_toolkit.logger import Logger
+
+
+class FunctionalBatchNorm(common.BaseSubstitution):
+    """
+    Replace functional batch_norm with BatchNorm2d.
+    """
+
+    def __init__(self):
+        """
+        Matches: functional batch_norm
+        """
+        bn_node = NodeOperationMatcher(F.batch_norm)
+        super().__init__(matcher_instance=bn_node)
+
+    def get_attributes_from_inputs(self, graph: Graph, node: BaseNode) -> dict:
+        input_nodes = graph.get_prev_nodes(node, sink_index_sorted=True)
+
+        if len(input_nodes) == 5:
+            return {
+                MOVING_MEAN: list(input_nodes[1].weights.values())[0],
+                MOVING_VARIANCE: list(input_nodes[2].weights.values())[0],
+                GAMMA: list(input_nodes[3].weights.values())[0],
+                BETA: list(input_nodes[4].weights.values())[0]
+            }
+        else:
+            Logger.warning(f'functional batch_norm is only folded in the 5 inputs case (input, mean, var, gamma, beta),'
+                           f'got {len(input_nodes)}')
+            return {}
+
+    def substitute(self,
+                   graph: Graph,
+                   node: BaseNode) -> Graph:
+        """
+        Substitute functional.batch_norm and its inputs with BatchNorm2d.
+        Args:
+            graph: Graph we apply the substitution on.
+            node: node that match the pattern in the substitution init.
+
+        Returns:
+            Graph after applying the substitution.
+        """
+        # if the input is not a 4D tensor, we can't substitute it with BatchNorm2d
+        if len(node.input_shape[0]) != 4:
+            return graph
+        out_channels = node.output_shape[0][1]
+
+        bn_node_weights = self.get_attributes_from_inputs(graph, node)
+        if not bn_node_weights:
+            return graph
+        new_batchnorm2d = BaseNode(name=node.name + '_into_BatchNorm2d',
+                                   framework_attr={NUM_FEATURES: out_channels,
+                                                   EPSILON: EPSILON_VAL,
+                                                   MOMENTUM: MOMENTUM_VAL},
+                                   input_shape=node.output_shape,
+                                   output_shape=node.output_shape,
+                                   weights=bn_node_weights,
+                                   layer_class=nn.BatchNorm2d)
+
+        num_nodes_before_substitution = len(graph.nodes)
+        num_edges_before_substitution = len(graph.edges)
+
+        batch_norm_consts = graph.get_prev_nodes(node)[1:]
+        for const in batch_norm_consts:
+            graph.remove_edge(const, node)
+            graph.remove_node(const)
+
+        graph.replace_node(node, new_batchnorm2d)
+
+        assert num_nodes_before_substitution - len(graph.nodes) == len(batch_norm_consts)
+        assert num_edges_before_substitution - len(graph.edges) == len(batch_norm_consts)
+
+        return graph

model_compression_toolkit/core/pytorch/pytorch_implementation.py

@@ -48,6 +48,8 @@ from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.ba
     pytorch_batchnorm_reconstruction
 from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.batchnorm_refusing import \
     pytorch_batchnorm_refusing
+from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.functional_batch_norm import \
+    FunctionalBatchNorm
 from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.linear_collapsing import \
     pytorch_linear_collapsing
 from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.multi_head_attention_decomposition \
@@ -243,7 +245,8 @@ class PytorchImplementation(FrameworkImplementation):
         return [ReshapeWithStaticShapes(),
                 MultiHeadAttentionDecomposition(),
                 PermuteCallMethod(),
-                ConstantHolderConv(fw_info)]
+                ConstantHolderConv(fw_info),
+                FunctionalBatchNorm()]
 
     def get_substitutions_pre_statistics_collection(self,
                                                     quant_config: QuantizationConfig