Modify DDP GPU test setting

examples/_test_requirements.txt

@@ -1 +1,2 @@
-scikit-learn
+scikit-learn
+jupyter

examples/uci/analyze.py

@@ -6,13 +6,12 @@
 
 import torch
 import torch.nn.functional as F
-from analyzer import Analyzer, prepare_model
-from arguments import FactorArguments, ScoreArguments
-from task import Task
 from torch import nn
-from torch.profiler import ProfilerActivity, profile, record_function
 
 from examples.uci.pipeline import construct_regression_mlp, get_regression_dataset
+from kronfluence.analyzer import Analyzer, prepare_model
+from kronfluence.arguments import FactorArguments
+from kronfluence.task import Task
 
 BATCH_DTYPE = Tuple[torch.Tensor, torch.Tensor]
 
@@ -32,31 +31,18 @@ def parse_args():
         default="./data",
         help="A folder containing the UCI regression dataset.",
     )
-
-    parser.add_argument(
-        "--factor_strategy",
-        type=str,
-        default="ekfac",
-        help="Strategy to compute preconditioning factors.",
-    )
-    parser.add_argument(
-        "--batch_size",
-        type=int,
-        default=256,
-        help="Batch size for compute factors and scores.",
-    )
     parser.add_argument(
-        "--analysis_name",
+        "--checkpoint_dir",
         type=str,
-        default="uci",
-        help="Name of the influence analysis.",
+        default="./checkpoints",
+        help="A path to store the final checkpoint.",
     )
 
     parser.add_argument(
-        "--checkpoint_dir",
+        "--factor_strategy",
         type=str,
-        default="./checkpoints",
-        help="A path to store the final checkpoint.",
+        default="ekfac",
+        help="Strategy to compute preconditioning factors.",
     )
 
     args = parser.parse_args()
@@ -110,7 +96,7 @@ def main():
     model = prepare_model(model, task)
 
     analyzer = Analyzer(
-        analysis_name=args.analysis_name,
+        analysis_name=args.dataset_name,
         model=model,
         task=task,
         cpu=True,
@@ -131,8 +117,7 @@ def main():
         factors_name=args.factor_strategy,
         query_dataset=eval_dataset,
         train_dataset=train_dataset,
-        per_device_query_batch_size=16,
-        per_device_train_batch_size=8,
+        per_device_query_batch_size=len(eval_dataset),
         overwrite_output_dir=True,
     )
     logging.info(f"Scores: {scores}")

examples/uci/tutorial.ipynb

@@ -2,15 +2,104 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "initial_id",
    "metadata": {
-    "collapsed": true
+    "collapsed": true,
+    "ExecuteTime": {
+     "end_time": "2024-03-12T10:46:20.005159Z",
+     "start_time": "2024-03-12T10:46:19.995640Z"
+    }
    },
    "outputs": [],
    "source": [
-    ""
+    "import kronfluence"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "outputs": [],
+   "source": [
+    "from examples.uci.train import train, evaluate"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-03-12T10:47:03.741170Z",
+     "start_time": "2024-03-12T10:47:02.235222Z"
+    }
+   },
+   "id": "4e56f0f1d6e34e62"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "outputs": [],
+   "source": [
+    "from kronfluence.analyzer import Analyzer, prepare_model\n",
+    "from kronfluence.arguments import FactorArguments\n",
+    "from kronfluence.task import Task"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-03-12T10:55:57.054379Z",
+     "start_time": "2024-03-12T10:55:57.034172Z"
+    }
+   },
+   "id": "6dc3ab20b6cb4050"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "outputs": [],
+   "source": [
+    "from examples.uci.pipeline import construct_regression_mlp, get_regression_dataset"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-03-12T11:40:46.724609Z",
+     "start_time": "2024-03-12T11:40:46.722860Z"
+    }
+   },
+   "id": "f3ed29a0d098c6dd"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "dataset_name = \"concrete\"\n",
+    "dataset_dir = \"./data\"\n",
+    "train_batch_size = 32\n",
+    "num_train_epochs = 40\n"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "cd2af4deeea3afd7"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "train_dataset = get_regression_dataset(data_name=dataset_name, split=\"train\", dataset_dir=dataset_dir)\n",
+    "\n",
+    "model = train(\n",
+    "    dataset=train_dataset,\n",
+    "    batch_size=args.train_batch_size,\n",
+    "    num_train_epochs=args.num_train_epochs,\n",
+    "    learning_rate=args.learning_rate,\n",
+    "    weight_decay=args.weight_decay,\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "c75658f17d06a7ab"
   }
  ],
  "metadata": {

kronfluence/analyzer.py

@@ -1,8 +1,6 @@
 from typing import Optional
 
 from accelerate.utils import extract_model_from_parallel
-from kronfluence.module.constants import FACTOR_TYPE
-
 from factor.config import FactorConfig
 from safetensors.torch import save_file
 from torch import nn
@@ -13,6 +11,7 @@
 from kronfluence.computer.eigen_computer import EigenComputer
 from kronfluence.computer.pairwise_score_computer import PairwiseScoreComputer
 from kronfluence.computer.self_score_computer import SelfScoreComputer
+from kronfluence.module.constants import FACTOR_TYPE
 from kronfluence.module.utils import wrap_tracked_modules
 from kronfluence.task import Task
 from kronfluence.utils.dataset import DataLoaderKwargs

kronfluence/score/pairwise.py

@@ -206,7 +206,7 @@ def compute_pairwise_scores_with_loaders(
     """Computes pairwise influence scores for a given model and task.
 
     Args:
-        loaded_factors (FACTOR_TYPE, optional):
+        loaded_factors (FACTOR_TYPE):
             The factor results to load from, before computing the pairwise scores.
         model (nn.Module):
             The model that pairwise influence scores will be computed.

kronfluence/score/self.py

@@ -93,7 +93,7 @@ def compute_self_scores_with_loaders(
     """Computes self-influence scores for a given model and task.
 
     Args:
-        loaded_factors (FACTOR_TYPE, optional):
+        loaded_factors (FACTOR_TYPE):
             The factor results to load from, before computing the self-influence scores.
         model (nn.Module):
             The model that self-influence scores will be computed.

tests/gpu_tests/compile_test.py

@@ -2,20 +2,21 @@
 import unittest
 
 import torch
-from analyzer import Analyzer, prepare_model
-from arguments import FactorArguments, ScoreArguments
-from module.constants import (
+from torch.utils import data
+
+from kronfluence.analyzer import Analyzer, prepare_model
+from kronfluence.arguments import FactorArguments, ScoreArguments
+from kronfluence.module.constants import (
     ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
     LAMBDA_FACTOR_NAMES,
 )
-from torch.utils import data
-
 from tests.gpu_tests.pipeline import (
     ClassificationTask,
     construct_mnist_mlp,
     get_mnist_dataset,
 )
+from tests.gpu_tests.prepare_tests import QUERY_INDICES, TRAIN_INDICES
 from tests.utils import check_tensor_dict_equivalence
 
 logging.basicConfig(level=logging.DEBUG)
@@ -33,9 +34,9 @@ def setUpClass(cls) -> None:
         cls.model = cls.model.double()
 
         cls.train_dataset = get_mnist_dataset(split="train", data_path="data")
-        cls.train_dataset = data.Subset(cls.train_dataset, indices=list(range(200)))
+        cls.train_dataset = data.Subset(cls.train_dataset, indices=list(range(TRAIN_INDICES)))
         cls.eval_dataset = get_mnist_dataset(split="valid", data_path="data")
-        cls.eval_dataset = data.Subset(cls.eval_dataset, indices=list(range(100)))
+        cls.eval_dataset = data.Subset(cls.eval_dataset, indices=list(range(QUERY_INDICES)))
 
         cls.task = ClassificationTask()
         cls.model = prepare_model(cls.model, cls.task)
@@ -60,7 +61,7 @@ def test_covariance_matrices(self) -> None:
             factors_name=NEW_FACTOR_NAME,
             dataset=self.train_dataset,
             factor_args=factor_args,
-            per_device_batch_size=16,
+            per_device_batch_size=512,
             overwrite_output_dir=True,
         )
         new_covariance_factors = self.analyzer.load_covariance_matrices(factors_name=NEW_FACTOR_NAME)
@@ -83,7 +84,7 @@ def test_lambda_matrices(self):
             factors_name=NEW_FACTOR_NAME,
             dataset=self.train_dataset,
             factor_args=factor_args,
-            per_device_batch_size=16,
+            per_device_batch_size=512,
             overwrite_output_dir=True,
             load_from_factors_name=OLD_FACTOR_NAME,
         )
@@ -114,10 +115,10 @@ def test_pairwise_scores(self) -> None:
             factors_name=OLD_FACTOR_NAME,
             query_dataset=self.eval_dataset,
             train_dataset=self.train_dataset,
-            train_indices=list(range(42)),
-            query_indices=list(range(23)),
-            per_device_query_batch_size=2,
-            per_device_train_batch_size=4,
+            train_indices=list(range(TRAIN_INDICES)),
+            query_indices=list(range(QUERY_INDICES)),
+            per_device_query_batch_size=12,
+            per_device_train_batch_size=512,
             score_args=score_args,
             overwrite_output_dir=True,
         )

tests/gpu_tests/cpu_test.py

@@ -2,15 +2,15 @@
 import unittest
 
 import torch
-from analyzer import Analyzer, prepare_model
-from arguments import FactorArguments, ScoreArguments
-from module.constants import (
+from torch.utils import data
+
+from kronfluence.analyzer import Analyzer, prepare_model
+from kronfluence.arguments import FactorArguments, ScoreArguments
+from kronfluence.module.constants import (
     ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
     LAMBDA_FACTOR_NAMES,
 )
-from torch.utils import data
-
 from tests.gpu_tests.pipeline import (
     ClassificationTask,
     construct_mnist_mlp,