Remove debug lines

pomonam · Jul 8, 2024 · 1dc7d85 · 1dc7d85
1 parent 3dce62d
commit 1dc7d85
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 63 deletions.
diff --git a/kronfluence/module/tracker/factor.py b/kronfluence/module/tracker/factor.py
@@ -255,7 +255,6 @@ def forward_hook(module: nn.Module, inputs: Tuple[torch.Tensor], outputs: torch.
 
         @torch.no_grad()
         def backward_hook(output_gradient: torch.Tensor) -> None:
-            print("Normal")
             if self.cached_activations is None:
                 self._raise_cache_not_found_exception()
             handle = self.cached_hooks.pop()
@@ -274,7 +273,6 @@ def backward_hook(output_gradient: torch.Tensor) -> None:
 
         @torch.no_grad()
         def shared_backward_hook(output_gradient: torch.Tensor) -> None:
-            print("Shared processed")
             handle = self.cached_hooks.pop()
             handle.remove()
             output_gradient = self._preprocess_gradient(

diff --git a/tests/gpu_tests/compile_test.py b/tests/gpu_tests/compile_test.py
@@ -9,6 +9,7 @@
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments, ScoreArguments
 from kronfluence.utils.common.factor_arguments import pytest_factor_arguments
+from kronfluence.utils.common.score_arguments import pytest_score_arguments
 from kronfluence.utils.constants import (
     ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
@@ -126,67 +127,58 @@ def test_lambda_shared_matrices(self) -> None:
                 rtol=RTOL,
             )
 
-    # def test_pairwise_scores(self) -> None:
-    #     pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=OLD_SCORE_NAME)
-    #
-    #     score_args = ScoreArguments(
-    #         score_dtype=torch.float64,
-    #         per_sample_gradient_dtype=torch.float64,
-    #         precondition_dtype=torch.float64,
-    #     )
-    #     self.analyzer.compute_pairwise_scores(
-    #         scores_name=NEW_SCORE_NAME,
-    #         factors_name=OLD_FACTOR_NAME,
-    #         query_dataset=self.eval_dataset,
-    #         train_dataset=self.train_dataset,
-    #         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,
-    #     )
-    #     new_pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=NEW_SCORE_NAME)
-    #
-    #     print(f"Previous score: {pairwise_scores[ALL_MODULE_NAME][10]}")
-    #     print(f"Previous shape: {pairwise_scores[ALL_MODULE_NAME].shape}")
-    #     print(f"New score: {new_pairwise_scores[ALL_MODULE_NAME][10]}")
-    #     print(f"New shape: {new_pairwise_scores[ALL_MODULE_NAME].shape}")
-    #     assert check_tensor_dict_equivalence(
-    #         pairwise_scores,
-    #         new_pairwise_scores,
-    #         atol=1e-5,
-    #         rtol=1e-3,
-    #     )
-    #
-    # def test_self_scores(self) -> None:
-    #     score_args = ScoreArguments(
-    #         score_dtype=torch.float64,
-    #         per_sample_gradient_dtype=torch.float64,
-    #         precondition_dtype=torch.float64,
-    #     )
-    #     self.analyzer.compute_self_scores(
-    #         scores_name=NEW_SCORE_NAME,
-    #         factors_name=OLD_FACTOR_NAME,
-    #         train_dataset=self.train_dataset,
-    #         train_indices=list(range(TRAIN_INDICES)),
-    #         per_device_train_batch_size=512,
-    #         score_args=score_args,
-    #         overwrite_output_dir=True,
-    #     )
-    #     new_self_scores = self.analyzer.load_self_scores(scores_name=NEW_SCORE_NAME)
-    #
-    #     self_scores = self.analyzer.load_self_scores(scores_name=OLD_SCORE_NAME)
-    #     print(f"Previous score: {self_scores[ALL_MODULE_NAME]}")
-    #     print(f"Previous shape: {self_scores[ALL_MODULE_NAME].shape}")
-    #     print(f"New score: {new_self_scores[ALL_MODULE_NAME]}")
-    #     print(f"New shape: {new_self_scores[ALL_MODULE_NAME].shape}")
-    #     assert check_tensor_dict_equivalence(
-    #         self_scores,
-    #         new_self_scores,
-    #         atol=1e-5,
-    #         rtol=1e-3,
-    #     )
+    def test_pairwise_scores(self) -> None:
+        score_args = pytest_score_arguments()
+        self.analyzer.compute_pairwise_scores(
+            scores_name=NEW_SCORE_NAME,
+            factors_name=OLD_FACTOR_NAME,
+            query_dataset=self.eval_dataset,
+            train_dataset=self.train_dataset,
+            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,
+        )
+        new_pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=NEW_SCORE_NAME)
+        pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=OLD_SCORE_NAME)
+
+        print(f"Previous score: {pairwise_scores[ALL_MODULE_NAME][10]}")
+        print(f"Previous shape: {pairwise_scores[ALL_MODULE_NAME].shape}")
+        print(f"New score: {new_pairwise_scores[ALL_MODULE_NAME][10]}")
+        print(f"New shape: {new_pairwise_scores[ALL_MODULE_NAME].shape}")
+        assert check_tensor_dict_equivalence(
+            pairwise_scores,
+            new_pairwise_scores,
+            atol=ATOL,
+            rtol=RTOL,
+        )
+
+    def test_self_scores(self) -> None:
+        score_args = pytest_score_arguments()
+        self.analyzer.compute_self_scores(
+            scores_name=NEW_SCORE_NAME,
+            factors_name=OLD_FACTOR_NAME,
+            train_dataset=self.train_dataset,
+            train_indices=list(range(TRAIN_INDICES)),
+            per_device_train_batch_size=512,
+            score_args=score_args,
+            overwrite_output_dir=True,
+        )
+        new_self_scores = self.analyzer.load_self_scores(scores_name=NEW_SCORE_NAME)
+        self_scores = self.analyzer.load_self_scores(scores_name=OLD_SCORE_NAME)
+
+        print(f"Previous score: {self_scores[ALL_MODULE_NAME]}")
+        print(f"Previous shape: {self_scores[ALL_MODULE_NAME].shape}")
+        print(f"New score: {new_self_scores[ALL_MODULE_NAME]}")
+        print(f"New shape: {new_self_scores[ALL_MODULE_NAME].shape}")
+        assert check_tensor_dict_equivalence(
+            self_scores,
+            new_self_scores,
+            atol=ATOL,
+            rtol=RTOL,
+        )
 
 
 if __name__ == "__main__":