super-gradients/tests/unit_tests/phase_context_test.py

import unittest

from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
from super_gradients.training.utils.callbacks import PhaseContextTestCallback, Phase
from super_gradients import Trainer
from super_gradients.training.metrics import Accuracy, Top5
from super_gradients.training.models import ResNet18
import torch
from super_gradients.training.utils.utils import AverageMeter
from torchmetrics import MetricCollection


class PhaseContextTest(unittest.TestCase):
    def context_information_in_train_test(self):
        trainer = Trainer("context_information_in_train_test")

        net = ResNet18(num_classes=5, arch_params={})

        phase_callbacks = [
            PhaseContextTestCallback(Phase.TRAIN_BATCH_END),
            PhaseContextTestCallback(Phase.TRAIN_BATCH_STEP),
            PhaseContextTestCallback(Phase.TRAIN_EPOCH_END),
            PhaseContextTestCallback(Phase.VALIDATION_BATCH_END),
            PhaseContextTestCallback(Phase.VALIDATION_EPOCH_END),
        ]

        train_params = {
            "max_epochs": 2,
            "lr_updates": [1],
            "lr_decay_factor": 0.1,
            "lr_mode": "step",
            "lr_warmup_epochs": 0,
            "initial_lr": 0.1,
            "loss": "cross_entropy",
            "optimizer": "SGD",
            "criterion_params": {},
            "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
            "train_metrics_list": [Accuracy()],
            "valid_metrics_list": [Top5()],
            "metric_to_watch": "Top5",
            "greater_metric_to_watch_is_better": True,
            "phase_callbacks": phase_callbacks,
        }

        trainer.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader())
        context_callbacks = list(filter(lambda cb: isinstance(cb, PhaseContextTestCallback), trainer.phase_callbacks))

        # CHECK THAT PHASE CONTEXES HAVE THE EXACT INFORMATION THERY'RE SUPPOSE TO HOLD
        for phase_callback in context_callbacks:
            if phase_callback.phase in [Phase.TRAIN_BATCH_END, Phase.TRAIN_BATCH_STEP, Phase.VALIDATION_BATCH_END]:
                self.assertTrue(phase_callback.context.batch_idx == 0)
                self.assertTrue(phase_callback.context.criterion is not None)
                self.assertTrue(isinstance(phase_callback.context.inputs, torch.Tensor))
                self.assertTrue(isinstance(phase_callback.context.loss_avg_meter, AverageMeter))
                self.assertTrue(isinstance(phase_callback.context.loss_log_items, torch.Tensor))
                self.assertTrue(phase_callback.context.metrics_dict is None)
                self.assertTrue(isinstance(phase_callback.context.preds, torch.Tensor))
                self.assertTrue(isinstance(phase_callback.context.target, torch.Tensor))

                if phase_callback.phase == Phase.VALIDATION_BATCH_END:
                    self.assertTrue(phase_callback.context.epoch == 2)
                    self.assertTrue(
                        isinstance(phase_callback.context.metrics_compute_fn, MetricCollection) and hasattr(phase_callback.context.metrics_compute_fn, "Top5")
                    )

                else:
                    self.assertTrue(phase_callback.context.epoch == 1)
                    self.assertTrue(
                        isinstance(phase_callback.context.metrics_compute_fn, MetricCollection)
                        and hasattr(phase_callback.context.metrics_compute_fn, "Accuracy")
                    )

        if phase_callback.phase in [Phase.TRAIN_EPOCH_END, Phase.VALIDATION_EPOCH_END]:
            self.assertTrue(phase_callback.context.batch_idx is None)
            self.assertTrue(phase_callback.context.criterion is None)
            self.assertTrue(phase_callback.context.inputs is None)
            self.assertTrue(phase_callback.context.loss_log_items is None)
            self.assertTrue(phase_callback.context.metrics_compute_fn is None)
            self.assertTrue(phase_callback.context.optimizer is not None)
            self.assertTrue(phase_callback.context.preds is None)
            self.assertTrue(phase_callback.context.target is None)
            self.assertTrue(phase_callback.context.epoch == 1)

            # EPOCH END PHASES USE THE SAME CONTEXT, WHICH IS UPDATED- SO VALID METRICS DICT SHOULD BE PRESENT
            self.assertTrue(isinstance(phase_callback.context.metrics_dict, dict))
            self.assertTrue("Loss" in phase_callback.context.metrics_dict.keys())
            self.assertTrue("Top5" in phase_callback.context.metrics_dict.keys())


if __name__ == "__main__":
    unittest.main()