Merge pull request #477 from aai-institute/fix/hessian_avg

Fix wrong averaging in hessian computation

Author: schroedk

notebooks/influence_imagenet.ipynb

@@ -432,14 +432,14 @@
      "text": [
       "INFO:pydvl.utils.progress:Function 'CgInfluence.influences' is starting.\n",
       "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' is starting.\n",
-      "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' completed. Duration: 0.63 sec\n",
+      "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' completed. Duration: 0.56 sec\n",
       "INFO:pydvl.utils.progress:Function 'CgInfluence._solve_hvp' is starting.\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3be47fd7d6cb40768cc3f229a629fa8a",
+       "model_id": "ecd37d2cced945beac116b384b789510",
        "version_major": 2,
        "version_minor": 0
       },
@@ -454,10 +454,10 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "INFO:pydvl.utils.progress:Function 'CgInfluence._solve_hvp' completed. Duration: 14.08 sec\n",
+      "INFO:pydvl.utils.progress:Function 'CgInfluence._solve_hvp' completed. Duration: 12.07 sec\n",
       "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' is starting.\n",
-      "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' completed. Duration: 3.76 sec\n",
-      "INFO:pydvl.utils.progress:Function 'CgInfluence.influences' completed. Duration: 18.48 sec\n"
+      "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' completed. Duration: 3.52 sec\n",
+      "INFO:pydvl.utils.progress:Function 'CgInfluence.influences' completed. Duration: 16.16 sec\n"
      ]
     }
    ],
@@ -873,14 +873,14 @@
      "text": [
       "INFO:pydvl.utils.progress:Function 'CgInfluence.influences' is starting.\n",
       "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' is starting.\n",
-      "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' completed. Duration: 0.60 sec\n",
+      "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' completed. Duration: 0.56 sec\n",
       "INFO:pydvl.utils.progress:Function 'CgInfluence._solve_hvp' is starting.\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f6313097fbe84b998cfcaa047c74f5f0",
+       "model_id": "a39f22193e37440a914254e643fde993",
        "version_major": 2,
        "version_minor": 0
       },
@@ -895,10 +895,10 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "INFO:pydvl.utils.progress:Function 'CgInfluence._solve_hvp' completed. Duration: 13.58 sec\n",
+      "INFO:pydvl.utils.progress:Function 'CgInfluence._solve_hvp' completed. Duration: 12.10 sec\n",
       "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' is starting.\n",
-      "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' completed. Duration: 3.64 sec\n",
-      "INFO:pydvl.utils.progress:Function 'CgInfluence.influences' completed. Duration: 17.82 sec\n"
+      "INFO:pydvl.utils.progress:Function 'TorchInfluenceFunctionModel._loss_grad' completed. Duration: 3.53 sec\n",
+      "INFO:pydvl.utils.progress:Function 'CgInfluence.influences' completed. Duration: 16.19 sec\n"
      ]
     }
    ],

notebooks/influence_synthetic.ipynb

@@ -407,6 +407,7 @@ def hessian(
     flat_params = flatten_dimensions(params.values())
 
     if use_hessian_avg:
+        n_samples = 0
         hessian = to_model_device(
             torch.zeros((n_parameters, n_parameters), dtype=model_dtype), model
         )
@@ -418,11 +419,12 @@ def flat_input_batch_loss_function(
             return blf(align_with_model(p, model), t_x, t_y)
 
         for x, y in iter(data_loader):
-            hessian += torch.func.hessian(flat_input_batch_loss_function)(
+            n_samples += x.shape[0]
+            hessian += x.shape[0] * torch.func.hessian(flat_input_batch_loss_function)(
                 flat_params, to_model_device(x, model), to_model_device(y, model)
             )
 
-        hessian /= len(data_loader)
+        hessian /= n_samples
     else:
 
         def flat_input_empirical_loss(p: torch.Tensor):