remove unnecessary code from hessian computation

sony · Nov 22, 2023 · e1199ce · e1199ce
1 parent a4054cc
commit e1199ce
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/model_compression_toolkit/core/keras/hessian/activation_trace_hessian_calculator_keras.py b/model_compression_toolkit/core/keras/hessian/activation_trace_hessian_calculator_keras.py
@@ -121,8 +121,7 @@ def compute(self) -> List[float]:
                             # Compute the approximation per node's output
                             score_approx_per_output = []
                             for grad in gradients:
-                                grad = tf.reshape(grad, [grad.shape[0], -1])
-                                score_approx_per_output.append(tf.reduce_mean(tf.reduce_sum(tf.pow(grad, 2.0))))
+                                score_approx_per_output.append(tf.reduce_sum(tf.pow(grad, 2.0)))
 
                             # Free gradients
                             del grad

diff --git a/...l_compression_toolkit/core/pytorch/hessian/activation_trace_hessian_calculator_pytorch.py b/...l_compression_toolkit/core/pytorch/hessian/activation_trace_hessian_calculator_pytorch.py
@@ -120,8 +120,7 @@ def compute(self) -> List[float]:
                                                      requires_grad=True,
                                                      device=device))
                         break
-                    hess_v = torch.reshape(hess_v, [hess_v.shape[0], -1])
-                    hessian_trace_approx = torch.mean(torch.sum(torch.pow(hess_v, 2.0)))
+                    hessian_trace_approx = torch.sum(torch.pow(hess_v, 2.0))
 
                     # If the change to the mean Hessian approximation is insignificant we stop the calculation
                     if j > MIN_HESSIAN_ITER: