super-gradients/tests/unit_tests/all_architectures_test.py

import unittest
from super_gradients.training.models.all_architectures import ARCHITECTURES
from super_gradients.training.models.sg_module import SgModule
from super_gradients.training.utils.utils import HpmStruct
import torch


class AllArchitecturesTest(unittest.TestCase):

    def setUp(self):
        # contains all arch_params needed for initialization of all architectures
        self.all_arch_params = HpmStruct(**{'num_classes': 10,
                                            'width_mult': 1,
                                            'threshold': 1,
                                            'sml_net': torch.nn.Identity(),
                                            'big_net': torch.nn.Identity(),
                                            'dropout': 0})

    def test_architecture_is_sg_module(self):
        """
        Validate all models from all_architectures.py are SgModule
        """
        for arch_name in ARCHITECTURES:
            # skip custom constructors to keep all_arch_params as general as a possible
            if 'custom' in arch_name.lower() or 'nas' in arch_name.lower():
                continue
            self.assertTrue(isinstance(ARCHITECTURES[arch_name](arch_params=self.all_arch_params), SgModule))


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