super-gradients/tests/unit_tests/registry_test.py

import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchmetrics
from torch.nn.modules.loss import _Loss

from super_gradients.training.models.all_architectures import ARCHITECTURES
from super_gradients.training.metrics.all_metrics import METRICS
from super_gradients.training.losses.all_losses import LOSSES
from super_gradients.common.registry import register_model, register_metric, register_loss


class RegistryTest(unittest.TestCase):

    def setUp(self):

        @register_model('myconvnet')
        class MyConvNet(nn.Module):
            def __init__(self, num_classes):
                super().__init__()
                self.conv1 = nn.Conv2d(3, 6, 5)
                self.pool = nn.MaxPool2d(2, 2)
                self.conv2 = nn.Conv2d(6, 16, 5)
                self.fc1 = nn.Linear(16 * 5 * 5, 120)
                self.fc2 = nn.Linear(120, 84)
                self.fc3 = nn.Linear(84, num_classes)

            def forward(self, x):
                x = self.pool(F.relu(self.conv1(x)))
                x = self.pool(F.relu(self.conv2(x)))
                x = torch.flatten(x, 1)
                x = F.relu(self.fc1(x))
                x = F.relu(self.fc2(x))
                x = self.fc3(x)
                return x

        @register_model()
        def myconvnet_for_cifar10():
            return MyConvNet(num_classes=10)

        @register_metric('custom_accuracy')  # Will be registered as "custom_accuracy"
        class CustomAccuracy(torchmetrics.Accuracy):
            def update(self, preds: torch.Tensor, target: torch.Tensor):
                if target.shape == preds.shape:
                    target = target.argmax(1)  # Supports smooth labels
                    super().update(preds=preds.argmax(1), target=target)

        @register_loss("custom_rsquared_loss")
        class CustomRSquaredLoss(_Loss):
            def forward(self, output, target):
                criterion_mse = nn.MSELoss()
                return 1 - criterion_mse(output, target).item() / torch.var(target).item()

    def tearDown(self):
        ARCHITECTURES.pop('myconvnet', None)
        ARCHITECTURES.pop('myconvnet_for_cifar10', None)
        METRICS.pop('custom_accuracy', None)
        LOSSES.pop('custom_rsquared_loss', None)

    def test_cls_is_registered(self):
        assert ARCHITECTURES['myconvnet']
        assert METRICS['custom_accuracy']
        assert LOSSES['custom_rsquared_loss']

    def test_fn_is_registered(self):
        assert ARCHITECTURES['myconvnet_for_cifar10']

    def test_is_instantiable(self):
        assert ARCHITECTURES['myconvnet_for_cifar10']()
        assert ARCHITECTURES['myconvnet'](num_classes=10)
        assert METRICS['custom_accuracy']()
        assert LOSSES['custom_rsquared_loss']()

    def test_model_outputs(self):
        torch.manual_seed(0)
        model_1 = ARCHITECTURES['myconvnet_for_cifar10']()
        torch.manual_seed(0)
        model_2 = ARCHITECTURES['myconvnet'](num_classes=10)
        dummy_input = torch.randn(1, 3, 32, 32, requires_grad=False)
        x = model_1(dummy_input)
        y = model_2(dummy_input)
        assert torch.equal(x, y)

    def test_existing_key(self):
        with self.assertRaises(Exception):
            @register_model()
            def myconvnet_for_cifar10():
                return


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