added test for validation forward wrapper

src/super_gradients/training/sg_trainer/sg_trainer.py

@@ -2289,8 +2289,8 @@ def evaluate(
                     else:
                         self.phase_callback_handler.on_test_batch_start(context)
 
-                    if self.self.validation_forward_wrapper is not None:
-                        output = self.validation_forward_wrapper(inputs, self.net)
+                    if self.validation_forward_wrapper is not None:
+                        output = self.validation_forward_wrapper(inputs=inputs, model=self.net)
                     else:
                         output = self.net(inputs)
 

tests/deci_core_unit_test_suite_runner.py

@@ -30,6 +30,7 @@
     TestClassificationAdapter,
     TestDetectionAdapter,
     TestSegmentationAdapter,
+    TestForwardWrapper,
 )
 from tests.end_to_end_tests import TestTrainer
 from tests.unit_tests.depth_estimation_dataset_test import DepthEstimationDatasetTest
@@ -192,6 +193,7 @@ def _add_modules_to_unit_tests_suite(self):
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(ClassBalancerTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(ClassBalancedSamplerTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestSegmentationModelExport))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestForwardWrapper))
 
     def _add_modules_to_end_to_end_tests_suite(self):
         """

tests/unit_tests/__init__.py

@@ -2,6 +2,7 @@
 from tests.unit_tests.crash_tips_test import CrashTipTest
 from tests.unit_tests.double_training_test import CallTrainTwiceTest
 from tests.unit_tests.factories_test import FactoriesTest
+from tests.unit_tests.forward_wrapper_test import TestForwardWrapper
 from tests.unit_tests.optimizer_params_override_test import TrainOptimizerParamsOverride
 from tests.unit_tests.resume_training_test import ResumeTrainingTest
 from tests.unit_tests.strictload_enum_test import StrictLoadEnumTest
@@ -63,4 +64,5 @@
     "TestClassificationAdapter",
     "TestDetectionAdapter",
     "TestSegmentationAdapter",
+    "TestForwardWrapper",
 ]

tests/unit_tests/forward_wrapper_test.py

@@ -0,0 +1,79 @@
+import unittest
+
+from super_gradients.training import Trainer
+from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
+from super_gradients.training.metrics import Accuracy
+from super_gradients.training.models import LeNet
+from super_gradients.training.utils.callbacks import PhaseContext, Callback
+import torch
+
+
+class OutputsCollectorCallback(Callback):
+    def __init__(self):
+        self.validation_outputs = []
+        self.train_outputs = []
+
+    def on_validation_batch_end(self, context: PhaseContext) -> None:
+        self.validation_outputs.append(context.preds)
+
+    def on_train_batch_end(self, context: PhaseContext) -> None:
+        self.train_outputs.append(context.preds)
+
+
+class DummyForwardWrapper:
+    def __call__(self, inputs: torch.Tensor, model: torch.nn.Module):
+        return torch.ones_like(model(inputs))
+
+
+def compare_tensor_lists(list1, list2):
+    if len(list1) != len(list2):
+        return False
+
+    # Move tensors to CPU
+    list1 = [t.cpu() for t in list1]
+    list2 = [t.cpu() for t in list2]
+
+    for tensor1, tensor2 in zip(list1, list2):
+        if not torch.all(torch.eq(tensor1, tensor2)):
+            return False
+    return True
+
+
+class TestForwardWrapper(unittest.TestCase):
+    def test_train_with_validation_forward_wrapper(self):
+        # Define Model
+        net = LeNet()
+        trainer = Trainer("test_train_with_validation_forward_wrapper")
+        output_collector = OutputsCollectorCallback()
+        validation_forward_wrapper = DummyForwardWrapper()
+        train_params = {
+            "max_epochs": 1,
+            "initial_lr": 0.1,
+            "loss": "CrossEntropyLoss",
+            "optimizer": "SGD",
+            "criterion_params": {},
+            "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
+            "train_metrics_list": [Accuracy()],
+            "valid_metrics_list": [Accuracy()],
+            "metric_to_watch": "Accuracy",
+            "greater_metric_to_watch_is_better": True,
+            "ema": False,
+            "phase_callbacks": [output_collector],
+            "warmup_mode": "LinearEpochLRWarmup",
+            "validation_forward_wrapper": validation_forward_wrapper,
+            "average_best_models": False,
+        }
+
+        expected_outputs = [torch.ones(4, 10)]
+        trainer.train(
+            model=net,
+            training_params=train_params,
+            train_loader=classification_test_dataloader(batch_size=4),
+            valid_loader=classification_test_dataloader(batch_size=4),
+        )
+        self.assertTrue(compare_tensor_lists(expected_outputs, output_collector.validation_outputs))
+        self.assertFalse(compare_tensor_lists(expected_outputs, output_collector.train_outputs))
+
+
+if __name__ == "__main__":
+    unittest.main()