super-gradients/tests/unit_tests/yolox_unit_test.py

import unittest

import torch

from super_gradients.training.losses import YoloXDetectionLoss, YoloXFastDetectionLoss
from super_gradients.training.models.detection_models.yolox import YoloX_N, YoloX_T, YoloX_S, YoloX_M, YoloX_L, YoloX_X
from super_gradients.training.utils.detection_utils import DetectionCollateFN
from super_gradients.training.utils.utils import HpmStruct


class TestYOLOX(unittest.TestCase):
    def setUp(self) -> None:
        self.arch_params = HpmStruct(num_classes=10)
        self.yolo_classes = [YoloX_N, YoloX_T, YoloX_S, YoloX_M, YoloX_L, YoloX_X]

    def test_yolox_creation(self):
        """
        test_yolox_creation - Tests the creation of the models
            :return:
        """
        dummy_input = torch.randn(1, 3, 320, 320)

        with torch.no_grad():

            for yolo_cls in self.yolo_classes:
                yolo_model = yolo_cls(self.arch_params)
                # THIS SHOULD RUN THE FORWARD ONCE
                yolo_model.eval()
                output_standard = yolo_model(dummy_input)
                self.assertIsNotNone(output_standard)

                # THIS SHOULD RUN A TRAINING FORWARD
                yolo_model.train()
                output_train = yolo_model(dummy_input)

                self.assertIsNotNone(output_train)

                # THIS SHOULD RUN THE FORWARD AUGMENT
                yolo_model.eval()
                yolo_model.augmented_inference = True
                output_augment = yolo_model(dummy_input)
                self.assertIsNotNone(output_augment)

    def test_yolox_loss(self):
        samples = [
            (torch.zeros((3, 256, 256)), torch.zeros((100, 5))),
            (torch.zeros((3, 256, 256)), torch.zeros((100, 5))),
            (torch.zeros((3, 256, 256)), torch.zeros((100, 5))),
            (torch.zeros((3, 256, 256)), torch.zeros((100, 5))),
            (torch.zeros((3, 256, 256)), torch.zeros((100, 5))),
        ]
        collate = DetectionCollateFN()
        _, targets = collate(samples)

        predictions = [
            torch.randn((5, 1, 256 // 8, 256 // 8, 4 + 1 + 10)),
            torch.randn((5, 1, 256 // 16, 256 // 16, 4 + 1 + 10)),
            torch.randn((5, 1, 256 // 32, 256 // 32, 4 + 1 + 10)),
        ]

        for loss in [
            YoloXDetectionLoss(strides=[8, 16, 32], num_classes=10, use_l1=True, iou_type="giou"),
            YoloXDetectionLoss(strides=[8, 16, 32], num_classes=10, use_l1=True, iou_type="iou"),
            YoloXDetectionLoss(strides=[8, 16, 32], num_classes=10, use_l1=False),
            YoloXFastDetectionLoss(strides=[8, 16, 32], num_classes=10, use_l1=True),
            YoloXFastDetectionLoss(strides=[8, 16, 32], num_classes=10, use_l1=False),
        ]:
            result = loss(predictions, targets)
            print(result)


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