replace normal distribution for random inputs with 1s and -1s

model_compression_toolkit/core/keras/hessian/activation_hessian_scores_calculator_keras.py

@@ -104,8 +104,8 @@ def compute(self) -> List[np.ndarray]:
             # Loop through each interest point activation tensor
             prev_mean_results = None
             for j in tqdm(range(self.num_iterations_for_approximation)):  # Approximation iterations
-                # Getting a random vector with normal distribution
-                v = tf.random.normal(shape=output.shape, dtype=output.dtype)
+                # Generate random tensor of 1s and -1s
+                v = self._generate_random_vectors_batch(output.shape)
                 f_v = tf.reduce_sum(v * output)
                 for i, ipt in enumerate(target_activation_tensors):  # Per Interest point activation tensor
                     interest_point_scores = []  # List to store scores for each interest point

model_compression_toolkit/core/keras/hessian/hessian_scores_calculator_keras.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from tensorflow import TensorShape
 
 from model_compression_toolkit.core.common.hessian.hessian_scores_calculator import HessianScoresCalculator
 
@@ -77,3 +78,19 @@ def _concat_tensors(self, tensors_to_concate: Union[tf.Tensor, List[tf.Tensor]])
                 "Unable to concatenate tensors for gradient calculation due to mismatched shapes along the first axis.") # pragma: no cover
 
         return tf.concat(_r_tensors, axis=1)
+
+    def _generate_random_vectors_batch(self, shape: TensorShape) -> tf.Tensor:
+        """
+        Generate a batch of random vectors for Hutchinson estimation using Rademacher distribution.
+
+        Args:
+            shape: target shape.
+
+        Returns:
+            Random tensor.
+        """
+        v = tf.random.uniform(shape=shape, minval=0, maxval=2, dtype=tf.int32)
+        v = tf.where(v == 0, -1, 1)
+        v = tf.cast(v, tf.float32)
+        return v
+

model_compression_toolkit/core/keras/hessian/weights_hessian_scores_calculator_keras.py

@@ -89,8 +89,7 @@ def compute(self) -> List[np.ndarray]:
             prev_mean_results = None
             tensors_original_shape = []
             for j in tqdm(range(self.num_iterations_for_approximation)):  # Approximation iterations
-                # Getting a random vector with normal distribution and the same shape as the model output
-                v = tf.random.normal(shape=output.shape)
+                v = self._generate_random_vectors_batch(output.shape)
                 f_v = tf.reduce_sum(v * output)
 
                 for i, ipt_node in enumerate(self.hessian_request.target_nodes):  # Per Interest point weights tensor