mct-nightly 2.2.0.20241022.507__py3-none-any.whl → 2.2.0.20241024.501__py3-none-any.whl

model_compression_toolkit/core/common/hessian/hessian_scores_request.py

@@ -12,10 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import List
+from typing import Iterable, Sequence, Optional, TYPE_CHECKING
+import dataclasses
 
 from enum import Enum
 
+if TYPE_CHECKING:    # pragma: no cover
+    from model_compression_toolkit.core.common import BaseNode
+
 
 class HessianMode(Enum):
     """
@@ -40,14 +44,7 @@ class HessianScoresGranularity(Enum):
     PER_TENSOR = 2
 
 
-class HessianEstimationDistribution(str, Enum):
-    """
-    Distribution for Hutchinson estimator random vector
-    """
-    GAUSSIAN = 'gaussian'
-    RADEMACHER = 'rademacher'
-
-
+@dataclasses.dataclass
 class HessianScoresRequest:
     """
     Request configuration for the Hessian-approximation scores.
@@ -55,36 +52,25 @@ class HessianScoresRequest:
     This class defines the parameters for the scores based on the Hessian matrix approximation.
     It specifies the mode (weights/activations), granularity (element/channel/tensor), and the target node.
 
-    Note: This does not compute scores using the actual Hessian matrix but an approximation.
+    Attributes:
+        mode: Mode of Hessian-approximation score (w.r.t weights or activations).
+        granularity: Granularity level for the approximation.
+        target_nodes: The node objects in the float graph for which the Hessian's approximation scores is targeted.
+        data_loader: Data loader to compute hessian approximations on. Should reflect the desired batch size for
+            the computation. Can be None if all hessians for the request are expected to be pre-computed previously.
+        n_samples: The number of samples to fetch hessian estimations for. If None, fetch hessians for a full pass
+            of the data loader.
     """
-
-    def __init__(self,
-                 mode: HessianMode,
-                 granularity: HessianScoresGranularity,
-                 target_nodes: List,
-                 distribution: HessianEstimationDistribution = HessianEstimationDistribution.GAUSSIAN):
-        """
-        Attributes:
-            mode (HessianMode): Mode of Hessian-approximation score (w.r.t weights or activations).
-            granularity (HessianScoresGranularity): Granularity level for the approximation.
-            target_nodes (List[BaseNode]): The node in the float graph for which the Hessian's approximation scores is targeted.
-        """
-
-        self.mode = mode  # w.r.t activations or weights
-        self.granularity = granularity  # per element, per layer, per channel
-        self.target_nodes = target_nodes
-        self.distribution = distribution
-
-    def __eq__(self, other):
-        # Checks if the other object is an instance of HessianScoresRequest
-        # and then checks if all attributes are equal.
-        return isinstance(other, HessianScoresRequest) and \
-               self.mode == other.mode and \
-               self.granularity == other.granularity and \
-               self.target_nodes == other.target_nodes and \
-               self.distribution == other.distribution
-
-    def __hash__(self):
-        # Computes the hash based on the attributes.
-        # The use of a tuple here ensures that the hash is influenced by all the attributes.
-        return hash((self.mode, self.granularity, tuple(self.target_nodes), self.distribution))
+    mode: HessianMode
+    granularity: HessianScoresGranularity
+    target_nodes: Sequence['BaseNode']
+    data_loader: Optional[Iterable]
+    n_samples: Optional[int]
+
+    def __post_init__(self):
+        if self.data_loader is None and self.n_samples is None:
+            raise ValueError('Data loader and the number of samples cannot both be None.')
+
+    def clone(self, **kwargs):
+        """ Create a clone with optional overrides """
+        return dataclasses.replace(self, **kwargs)

model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py

@@ -238,22 +238,20 @@ class SensitivityEvaluation:
         """
         # Create a request for Hessian approximation scores with specific configurations
         # (here we use per-tensor approximation of the Hessian's trace w.r.t the node's activations)
+        fw_dataloader = self.fw_impl.convert_data_gen_to_dataloader(self.representative_data_gen,
+                                                                    batch_size=self.quant_config.hessian_batch_size)
         hessian_info_request = HessianScoresRequest(mode=HessianMode.ACTIVATION,
                                                     granularity=HessianScoresGranularity.PER_TENSOR,
-                                                    target_nodes=self.interest_points)
+                                                    target_nodes=self.interest_points,
+                                                    data_loader=fw_dataloader,
+                                                    n_samples=self.quant_config.num_of_images)
 
         # Fetch the Hessian approximation scores for the current interest point
-        nodes_approximations = self.hessian_info_service.fetch_hessian(hessian_scores_request=hessian_info_request,
-                                                                       required_size=self.quant_config.num_of_images,
-                                                                       batch_size=self.quant_config.hessian_batch_size)
-
-        # Store the approximations for each node for each image
-        approx_by_image = [[nodes_approximations[j][image_idx]
-                            for j, _ in enumerate(self.interest_points)]
-                           for image_idx in range(self.quant_config.num_of_images)]
+        nodes_approximations = self.hessian_info_service.fetch_hessian(request=hessian_info_request)
+        approx_by_image = np.stack([nodes_approximations[n.name] for n in self.interest_points], axis=1)    # samples X nodes
 
         # Return the mean approximation value across all images for each interest point
-        return np.mean(np.stack(approx_by_image), axis=0)
+        return np.mean(approx_by_image, axis=0)
 
     def _configure_bitwidths_model(self,
                                    mp_model_configuration: List[int],

model_compression_toolkit/core/common/pruning/importance_metrics/lfh_importance_metric.py

@@ -120,22 +120,24 @@ class LFHImportanceMetric(BaseImportanceMetric):
         """
 
         # Initialize HessianInfoService for score computation.
+
         hessian_info_service = HessianInfoService(graph=self.float_graph,
-                                                  representative_dataset_gen=self.representative_data_gen,
                                                   fw_impl=self.fw_impl)
 
         # Fetch and process Hessian scores for output channels of entry nodes.
-        nodes_scores = []
+        data_loader = self.fw_impl.convert_data_gen_to_dataloader(self.representative_data_gen, batch_size=1)
+        nodes_scores = {}
         for node in entry_nodes:
-            _request = HessianScoresRequest(mode=HessianMode.WEIGHTS,
-                                            granularity=HessianScoresGranularity.PER_OUTPUT_CHANNEL,
-                                            target_nodes=[node])
-            _scores_for_node = hessian_info_service.fetch_hessian(_request,
-                                                                  required_size=self.pruning_config.num_score_approximations)
-            nodes_scores.append(_scores_for_node)
+            request = HessianScoresRequest(mode=HessianMode.WEIGHTS,
+                                           granularity=HessianScoresGranularity.PER_OUTPUT_CHANNEL,
+                                           target_nodes=[node],
+                                           data_loader=data_loader,
+                                           n_samples=self.pruning_config.num_score_approximations)
+            node_scores = hessian_info_service.fetch_hessian(request)
+            nodes_scores.update(node_scores)
 
         # Average and map scores to nodes.
-        self._entry_node_to_hessian_score = {node: np.mean(scores[0], axis=0).flatten() for node, scores in zip(entry_nodes, nodes_scores)}
+        self._entry_node_to_hessian_score = {node: np.mean(nodes_scores[node.name], axis=0).flatten() for node in entry_nodes}
 
         self._entry_node_count_oc_nparams = self._count_oc_nparams(entry_nodes=entry_nodes)
         _entry_node_l2_oc_norm = self._get_squaredl2norm(entry_nodes=entry_nodes)

model_compression_toolkit/core/common/quantization/quantization_params_generation/error_functions.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 from copy import deepcopy
-from typing import Tuple, Callable, List
+from typing import Tuple, Callable, List, Iterable, Optional
 import numpy as np
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
 from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianMode, HessianScoresGranularity, \
@@ -377,7 +377,8 @@ def _get_sliced_histogram(bins: np.ndarray,
 
 def _compute_hessian_for_hmse(node,
                               hessian_info_service: HessianInfoService,
-                              num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> List[List[np.ndarray]]:
+                              num_hessian_samples: int,
+                              dataloader: Optional[Iterable]) -> List[List[np.ndarray]]:
     """
     Compute and retrieve Hessian-based scores for using during HMSE error computation.
 
@@ -385,15 +386,18 @@ def _compute_hessian_for_hmse(node,
         node: The node to compute Hessian-based scores for.
         hessian_info_service: HessianInfoService object for retrieving Hessian-based scores.
         num_hessian_samples: Number of samples to approximate Hessian-based scores on.
+        dataloader: Data loader for computing Hessian-based scores. Can be None if hessians are expected to be
+          available, i.e. have been already computed previously.
 
     Returns: A list with computed Hessian-based scores tensors for the given node.
 
     """
     _request = HessianScoresRequest(mode=HessianMode.WEIGHTS,
                                     granularity=HessianScoresGranularity.PER_ELEMENT,
+                                    data_loader=dataloader,
+                                    n_samples=num_hessian_samples,
                                     target_nodes=[node])
-    _scores_for_node = hessian_info_service.fetch_hessian(_request,
-                                                          required_size=num_hessian_samples)
+    _scores_for_node = hessian_info_service.fetch_hessian(_request)
 
     return _scores_for_node
 
@@ -476,11 +480,11 @@ def get_threshold_selection_tensor_error_function(quantization_method: Quantizat
                                                                           per_channel=True)
 
     if quant_error_method == qc.QuantizationErrorMethod.HMSE:
-        node_hessian_scores = _compute_hessian_for_hmse(node, hessian_info_service, num_hessian_samples)
+        node_hessian_scores = _compute_hessian_for_hmse(node, hessian_info_service, num_hessian_samples, None)
         if len(node_hessian_scores) != 1:
             Logger.critical(f"Expecting single node Hessian score request to return a list of length 1, but got a list "
                             f"of length {len(node_hessian_scores)}.")
-        node_hessian_scores = np.sqrt(np.mean(node_hessian_scores[0], axis=0))
+        node_hessian_scores = np.sqrt(np.mean(node_hessian_scores[node.name], axis=0))
 
         return lambda x, y, threshold: _hmse_error_function_wrapper(x, y, norm=norm, axis=axis,
                                                                     hessian_scores=node_hessian_scores)

model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_computation.py

@@ -15,11 +15,12 @@
 import copy
 
 from tqdm import tqdm
-from typing import List
+from typing import List, Callable, Generator
 
 from model_compression_toolkit.constants import NUM_QPARAM_HESSIAN_SAMPLES
 from model_compression_toolkit.core import QuantizationErrorMethod
 from model_compression_toolkit.core.common import Graph, BaseNode
+from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.hessian import HessianInfoService, HessianScoresRequest, HessianMode, \
     HessianScoresGranularity
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_activations_computation \
@@ -55,26 +56,25 @@ def _collect_nodes_for_hmse(nodes_list: List[BaseNode], graph: Graph) -> List[Ba
 
 
 def calculate_quantization_params(graph: Graph,
-                                  nodes: List[BaseNode] = [],
-                                  specific_nodes: bool = False,
+                                  fw_impl: FrameworkImplementation,
+                                  repr_data_gen_fn: Callable[[], Generator],
+                                  nodes: List[BaseNode] = None,
                                   hessian_info_service: HessianInfoService = None,
                                   num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES):
     """
     For a graph, go over its nodes, compute quantization params (for both weights and activations according
     to the given framework info), and create and attach a NodeQuantizationConfig to each node (containing the
     computed params).
-    By default, the function goes over all nodes in the graph. However, the specific_nodes flag enables
-    to compute quantization params for specific nodes if the default behavior is unnecessary. For that,
-    a list of nodes should be passed as well.
+    By default, the function goes over all nodes in the graph. However, specific nodes can be passed
+    to compute quantization params only for them.
 
     Args:
-        groups of layers by how they should be quantized, etc.)
         graph: Graph to compute its nodes' thresholds.
+        fw_impl: FrameworkImplementation object.
+        repr_data_gen_fn: callable returning representative dataset generator.
         nodes: List of nodes to compute their thresholds instead of computing it for all nodes in the graph.
-        specific_nodes: Flag to compute thresholds for only specific nodes.
         hessian_info_service: HessianInfoService object for retrieving Hessian-based scores (used only with HMSE error method).
         num_hessian_samples: Number of samples to approximate Hessian-based scores on (used only with HMSE error method).
-
     """
 
     Logger.info(f"\nRunning quantization parameters search. "
@@ -82,18 +82,20 @@ def calculate_quantization_params(graph: Graph,
                 f"depending on the model size and the selected quantization methods.\n")
 
     # Create a list of nodes to compute their thresholds
-    nodes_list: List[BaseNode] = nodes if specific_nodes else graph.nodes()
+    nodes_list: List[BaseNode] = nodes or graph.nodes()
 
     # Collecting nodes that are configured to search weights quantization parameters using HMSE optimization
     # and computing required Hessian information to be used for HMSE parameters selection.
     # The Hessian scores are computed and stored in the hessian_info_service object.
     nodes_for_hmse = _collect_nodes_for_hmse(nodes_list, graph)
     if len(nodes_for_hmse) > 0:
-        hessian_info_service.fetch_hessian(HessianScoresRequest(mode=HessianMode.WEIGHTS,
-                                                                granularity=HessianScoresGranularity.PER_ELEMENT,
-                                                                target_nodes=nodes_for_hmse),
-                                           required_size=num_hessian_samples,
-                                           batch_size=1)
+        dataloader = fw_impl.convert_data_gen_to_dataloader(repr_data_gen_fn, batch_size=1)
+        request = HessianScoresRequest(mode=HessianMode.WEIGHTS,
+                                       granularity=HessianScoresGranularity.PER_ELEMENT,
+                                       data_loader=dataloader,
+                                       n_samples=num_hessian_samples,
+                                       target_nodes=nodes_for_hmse)
+        hessian_info_service.fetch_hessian(request)
 
     for n in tqdm(nodes_list, "Calculating quantization parameters"):  # iterate only nodes that we should compute their thresholds
         for candidate_qc in n.candidates_quantization_cfg:

model_compression_toolkit/core/keras/data_util.py

@@ -0,0 +1,67 @@
+# Copyright 2024 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 typing import Generator, Callable
+
+import tensorflow as tf
+
+from model_compression_toolkit.core.keras.tf_tensor_numpy import to_tf_tensor
+
+
+def flat_gen_fn(data_gen_fn: Callable[[], Generator]):
+    """
+    Convert data generator with arbitrary batch size to a flat (sample by sample) data generator.
+
+    Args:
+        data_gen_fn: input data generator factory. Generator is expected to yield lists of tensors.
+
+    Returns:
+        A factory for a flattened data generator.
+    """
+    def gen():
+        for inputs_batch in data_gen_fn():
+            for sample in zip(*inputs_batch):
+                yield to_tf_tensor(sample)
+    return gen
+
+
+# TODO in tf dataset and dataloader are combined within tf.data.Dataset. For advanced use cases such as gptq sla we
+#  need to separate dataset from dataloader similarly to torch data_util.
+class TFDatasetFromGenerator:
+    def __init__(self, data_gen, batch_size):
+        inputs = next(data_gen())
+        if not isinstance(inputs, list):
+            raise TypeError(f'Representative data generator is expected to generate a list of tensors, '
+                            f'got {type(inputs)}')  # pragma: no cover
+
+        self.orig_batch_size = inputs[0].shape[0]
+
+        output_signature = tuple([tf.TensorSpec(shape=t.shape[1:], dtype=t.dtype) for t in inputs])
+        dataset = tf.data.Dataset.from_generator(flat_gen_fn(data_gen), output_signature=output_signature)
+        self.dataset = dataset.batch(batch_size)
+        self._size = None
+
+    def __iter__(self):
+        return iter(self.dataset)
+
+    def __len__(self):
+        """ Returns the number of batches. """
+        if self._size is None:
+            self._num_batches = sum(1 for _ in self)
+        return self._num_batches
+
+
+def data_gen_to_dataloader(data_gen_fn: Callable[[], Generator], batch_size) -> TFDatasetFromGenerator:
+    """ Create DataLoader based on samples yielded by data_gen. """
+    return TFDatasetFromGenerator(data_gen_fn, batch_size)

model_compression_toolkit/core/keras/keras_implementation.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 from functools import partial
-from typing import List, Any, Tuple, Callable, Dict, Union
+from typing import List, Any, Tuple, Callable, Dict, Union, Generator
 
 import numpy as np
 import tensorflow as tf
@@ -23,6 +23,7 @@ from tensorflow.keras.models import Model
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianMode, HessianInfoService
+from model_compression_toolkit.core.keras.data_util import data_gen_to_dataloader
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.remove_identity import RemoveIdentity
 from model_compression_toolkit.core.keras.hessian.activation_hessian_scores_calculator_keras import \
     ActivationHessianScoresCalculatorKeras
@@ -628,3 +629,8 @@ class KerasImplementation(FrameworkImplementation):
                                         get_weights_quantizer_for_node,
                                         get_activations_quantizer_for_node,
                                         attribute_names)
+
+    @staticmethod
+    def convert_data_gen_to_dataloader(data_gen_fn: Callable[[], Generator], batch_size: int):
+        """ Create DataLoader based on samples yielded by data_gen. """
+        return data_gen_to_dataloader(data_gen_fn, batch_size=batch_size)

model_compression_toolkit/core/keras/tf_tensor_numpy.py

@@ -33,7 +33,7 @@ def to_tf_tensor(tensor):
     elif isinstance(tensor, list):
         return [to_tf_tensor(t) for t in tensor]
     elif isinstance(tensor, tuple):
-        return (to_tf_tensor(t) for t in tensor)
+        return tuple(to_tf_tensor(t) for t in tensor)
     elif isinstance(tensor, np.ndarray):
         return tf.convert_to_tensor(tensor.astype(np.float32))
     else:  # pragma: no cover

model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py

@@ -134,7 +134,7 @@ def _run_operation(n: BaseNode,
         input_tensors = n.insert_positional_weights_to_input_list(input_tensors)
         # convert inputs from positional weights (numpy arrays) to tensors. Must handle each element in the
         # list separately, because in FX the tensors are FX objects and fail to_torch_tensor
-        input_tensors = [to_torch_tensor(t, numpy_type=t.dtype) if isinstance(t, np.ndarray) else t
+        input_tensors = [to_torch_tensor(t, None) if isinstance(t, np.ndarray) else t
                          for t in input_tensors]
         _tensor_input_allocs = None
 

model_compression_toolkit/core/pytorch/data_util.py

@@ -0,0 +1,163 @@
+# Copyright 2024 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 typing import Generator, Callable, Sequence, Any
+
+import torch
+from torch.utils.data import IterableDataset, Dataset, DataLoader, default_collate
+
+
+def flat_gen_fn(data_gen_fn: Callable[[], Generator]):
+    """
+    Convert data generator with arbitrary batch size to a flat (sample by sample) data generator.
+
+    Args:
+        data_gen_fn: input data generator factory. Generator is expected to yield lists of tensors.
+
+    Returns:
+        A factory for a flattened data generator.
+    """
+    def gen():
+        for inputs_batch in data_gen_fn():
+            for sample in zip(*inputs_batch):
+                # convert to torch tensor but do not move to device yet (it will cause issues with num_workers > 0)
+                yield [torch.as_tensor(s) for s in sample]
+    return gen
+
+
+class IterableDatasetFromGenerator(IterableDataset):
+    """
+    PyTorch iterable dataset built from a data generator factory.
+    Each iteration over the dataset corresponds to one pass over a fresh instance of a data generator.
+    Therefore, if the data generator factory creates data generator instances that yield different samples,
+    this behavior is preserved.
+    """
+
+    def __init__(self, data_gen_fn: Callable[[], Generator]):
+        """
+        Args:
+            data_gen_fn: a factory for data generator that yields lists of tensors.
+        """
+        # validate one batch
+        test_batch = next(data_gen_fn())
+        if not isinstance(test_batch, list):
+            raise TypeError(f'Data generator is expected to yield a list of tensors, got {type(test_batch)}')
+        self.orig_batch_size = test_batch[0].shape[0]
+
+        self._size = None
+        self._gen_fn = flat_gen_fn(data_gen_fn)
+
+    def __iter__(self):
+        """ Return an iterator for the dataset. """
+        return self._gen_fn()
+
+    def __len__(self):
+        """ Get the length of the dataset. """
+        if self._size is None:
+            self._size = sum(1 for _ in self)
+        return self._size
+
+
+class FixedDatasetFromGenerator(Dataset):
+    """
+    Dataset containing a fixed number of samples (i.e. same samples are yielded in each epoch), retrieved from a
+    data generator.
+    Note that the samples are stored in memory.
+
+    Attributes:
+        orig_batch_size: the batch size of the input data generator (retrieved from the first batch).
+    """
+    def __init__(self, data_gen_fn: Callable[[], Generator], n_samples: int = None):
+        """
+        Args:
+            data_gen_fn: data generator factory.
+            n_samples: target size of the dataset. If None, use all samples yielded by the data generator in one pass.
+        """
+        test_batch = next(data_gen_fn())
+        if not isinstance(test_batch, list):
+            raise TypeError(f'Data generator is expected to yield a list of tensors, got {type(test_batch)}')
+        self.orig_batch_size = test_batch[0].shape[0]
+
+        samples = []
+        for batch in data_gen_fn():
+            # convert to torch tensor but do not move to device yet (it will cause issues with num_workers > 0)
+            batch = [torch.as_tensor(t) for t in batch]
+            samples.extend(zip(*batch))
+            if n_samples is not None and len(samples) >= n_samples:
+                samples = samples[:n_samples]
+                break
+
+        if n_samples is not None and len(samples) < n_samples:
+            raise ValueError(f'Not enough samples in the data generator to create a dataset with {n_samples}')
+        self.samples = samples
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        return list(self.samples[index])
+
+
+class FixedSampleInfoDataset(Dataset):
+    """
+    Dataset for samples augmented with additional info per sample.
+    Each element in the dataset is a tuple containing the sample and sample's additional info.
+    """
+    def __init__(self, samples: Sequence, *sample_info: Sequence):
+        """
+        Args:
+            samples: a sequence of input samples.
+            hessians: one or more sequences of samples complementary data of matching sizes.
+        """
+        if not all(len(info) == len(samples) for info in sample_info):
+            raise ValueError('Mismatch in the number of samples between samples and complementary data.')
+        self.samples = samples
+        self.sample_info = sample_info
+
+    def __getitem__(self, index):
+        return self.samples[index], *[info[index] for info in self.sample_info]
+
+    def __len__(self):
+        return len(self.samples)
+
+
+class IterableSampleWithConstInfoDataset(IterableDataset):
+    """
+    A convenience dataset that augments each sample with additional info shared by all samples.
+    """
+    def __init__(self, samples_dataset: Dataset, *info: Any):
+        """
+        Args:
+            samples_dataset: any dataset containing samples.
+            *sample_info: one or more static entities to augment each sample.
+        """
+        self.samples_dataset = samples_dataset
+        self.info = info
+
+    def __iter__(self):
+        for sample in self.samples_dataset:
+            yield sample, *self.info
+
+
+def get_collate_fn_with_extra_outputs(*extra_outputs: Any) -> Callable:
+    """ Collation function that adds const extra outputs to each batch. """
+    def f(batch):
+        return default_collate(batch) + list(extra_outputs)
+    return f
+
+
+def data_gen_to_dataloader(data_gen_fn: Callable[[], Generator], batch_size, **kwargs):
+    """ Create DataLoader based on samples yielded by data_gen. """
+    dataset = IterableDatasetFromGenerator(data_gen_fn)
+    return DataLoader(dataset, batch_size=batch_size, **kwargs)

model_compression_toolkit/core/pytorch/hessian/activation_hessian_scores_calculator_pytorch.py

@@ -15,20 +15,19 @@
 
 from typing import List
 
+import numpy as np
+import torch
 from torch import autograd
 from tqdm import tqdm
-import numpy as np
 
 from model_compression_toolkit.constants import MIN_HESSIAN_ITER, HESSIAN_COMP_TOLERANCE, HESSIAN_NUM_ITERATIONS
 from model_compression_toolkit.core.common import Graph
-from model_compression_toolkit.core.common.hessian import (HessianScoresRequest, HessianScoresGranularity,
-                                                           HessianEstimationDistribution)
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianScoresGranularity
 from model_compression_toolkit.core.pytorch.back2framework.float_model_builder import FloatPyTorchModelBuilder
 from model_compression_toolkit.core.pytorch.hessian.hessian_scores_calculator_pytorch import \
     HessianScoresCalculatorPytorch
 from model_compression_toolkit.core.pytorch.utils import torch_tensor_to_numpy
 from model_compression_toolkit.logger import Logger
-import torch
 
 
 class ActivationHessianScoresCalculatorPytorch(HessianScoresCalculatorPytorch):
@@ -86,36 +85,12 @@ class ActivationHessianScoresCalculatorPytorch(HessianScoresCalculatorPytorch):
         target_activation_tensors = outputs[:num_target_nodes]
         # Extract the model outputs
         output_tensors = outputs[num_target_nodes:]
-        device = output_tensors[0].device
 
         # Concat outputs
         # First, we need to unfold all outputs that are given as list, to extract the actual output tensors
         output = self.concat_tensors(output_tensors)
         return output, target_activation_tensors
 
-    def _generate_random_vectors_batch(self, shape: tuple, distribution: HessianEstimationDistribution,
-                                       device: torch.device) -> torch.Tensor:
-        """
-        Generate a batch of random vectors for Hutchinson estimation
-
-        Args:
-            shape: target shape
-            distribution: distribution to sample from
-            device: target device
-
-        Returns:
-            Random tensor
-        """
-        if distribution == HessianEstimationDistribution.GAUSSIAN:
-            return torch.randn(shape, device=device)
-
-        if distribution == HessianEstimationDistribution.RADEMACHER:
-            v = torch.randint(high=2, size=shape, device=device)
-            v[v == 0] = -1
-            return v
-
-        raise ValueError(f'Unknown distribution {distribution}')    # pragma: no cover
-
     def compute(self) -> List[np.ndarray]:
         """
         Compute the scores that are based on the approximation of the Hessian w.r.t the requested target nodes' activations.
@@ -142,8 +117,8 @@ class ActivationHessianScoresCalculatorPytorch(HessianScoresCalculatorPytorch):
                                       for _ in range(len(target_activation_tensors))]
         prev_mean_results = None
         for j in tqdm(range(self.num_iterations_for_approximation), "Hessian random iterations"):  # Approximation iterations
-            # Getting a random vector with normal distribution
-            v = self._generate_random_vectors_batch(output.shape, self.hessian_request.distribution, output.device)
+            # Getting a random vector
+            v = self._generate_random_vectors_batch(output.shape, output.device)
             f_v = torch.sum(v * output)
             for i, ipt_tensor in enumerate(target_activation_tensors):  # Per Interest point activation tensor
                 # Computing the hessian-approximation scores by getting the gradient of (output * v)
@@ -184,7 +159,7 @@ class ActivationHessianScoresCalculatorPytorch(HessianScoresCalculatorPytorch):
                                       for _ in range(len(target_activation_tensors))]
 
         for j in tqdm(range(self.num_iterations_for_approximation), "Hessian random iterations"):  # Approximation iterations
-            v = self._generate_random_vectors_batch(output.shape, self.hessian_request.distribution, output.device)
+            v = self._generate_random_vectors_batch(output.shape, output.device)
             f_v = torch.sum(v * output)
             for i, ipt_tensor in enumerate(target_activation_tensors):  # Per Interest point activation tensor
                 hess_v = autograd.grad(outputs=f_v,