@@ -5,6 +5,15 @@ from super_gradients.common.factories.metrics_factory import MetricsFactory
 from super_gradients.common.factories.optimizers_type_factory import OptimizersTypeFactory
 from super_gradients.common.factories.samplers_factory import SamplersFactory
 from super_gradients.common.factories.transforms_factory import TransformsFactory
+from super_gradients.common.factories.activations_type_factory import ActivationsTypeFactory
 
-
-__all__ = ["CallbacksFactory", "ListFactory", "LossesFactory", "MetricsFactory", "OptimizersTypeFactory", "SamplersFactory", "TransformsFactory"]
+__all__ = [
+    "CallbacksFactory",
+    "ListFactory",
+    "LossesFactory",
+    "MetricsFactory",
+    "OptimizersTypeFactory",
+    "SamplersFactory",
+    "TransformsFactory",
+    "ActivationsTypeFactory",
+]

from typing import Union, Type, Mapping

from super_gradients.common.factories.base_factory import AbstractFactory
from super_gradients.training.utils.activations_utils import get_builtin_activation_type


class ActivationsTypeFactory(AbstractFactory):
    """
    This is a special factory for getting a type of the activation function by name.
    This factory does not instantiate a module, but rather return the type to be instantiated via call method.
    """

    def get(self, conf: Union[str, Mapping]) -> Type:
        """
        Get a type.
           :param conf: a configuration
           if string - assumed to be a type name (not the real name, but a name defined in the Factory)
           a dictionary is not supported, since the actual instantiation takes place elsewhere

           If provided value is not one of the three above, the value will be returned as is
        """
        if isinstance(conf, str):
            return get_builtin_activation_type(conf)

        if isinstance(conf, Mapping):
            (type_name,) = list(conf.keys())
            type_args = conf[type_name]
            return get_builtin_activation_type(type_name, **type_args)

        raise RuntimeError(f"Unsupported conf param type {type(conf)}")

from functools import partial
from typing import Type, Union, Dict

import torch
from torch import nn


def get_builtin_activation_type(activation: Union[str, None], **kwargs) -> Type:
    """
    Returns activation class by its name from torch.nn namespace. This function support all modules available from
    torch.nn and also their lower-case aliases.
    On top of that, it supports a few aliaes: leaky_relu (LeakyReLU), swish (silu).

    >>> act_cls = get_activation_type("LeakyReLU", inplace=True, slope=0.01)
    >>> act = act_cls()

    Args:
        activation: Activation function name (E.g. ReLU). If None will return nn.Identity
        **kwargs: Extra arguments to pass to constructor during instantiation (E.g. inplace=True)

    Returns:
        Type of the activation function that is ready to be instantiated
    """

    if activation is None:
        activation_cls = nn.Identity
    else:
        lowercase_aliases: Dict[str, str] = dict((k.lower(), k) for k in torch.nn.__dict__.keys())

        # Register additional aliases
        lowercase_aliases["leaky_relu"] = "LeakyReLU"  # LeakyRelu in snake_case
        lowercase_aliases["swish"] = "SiLU"  # Swish shich is equivalent to SiLU
        lowercase_aliases["none"] = "Identity"

        if activation in lowercase_aliases:
            activation = lowercase_aliases[activation]

        if activation not in torch.nn.__dict__:
            raise KeyError(f"Requested activation function {activation} is not known")

        activation_cls = torch.nn.__dict__[activation]
        if len(kwargs):
            activation_cls = partial(activation_cls, **kwargs)

    return activation_cls

@@ -3,9 +3,12 @@ import unittest
 import torch
 
 from super_gradients import Trainer
+from super_gradients.common.decorators.factory_decorator import resolve_param
+from super_gradients.common.factories import ActivationsTypeFactory
 from super_gradients.training import models
 from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
 from super_gradients.training.metrics import Accuracy, Top5
+from torch import nn
 
 
 class FactoriesTest(unittest.TestCase):
@@ -36,6 +39,16 @@ class FactoriesTest(unittest.TestCase):
         self.assertIsInstance(trainer.valid_metrics.Top5, Top5)
         self.assertIsInstance(trainer.optimizer, torch.optim.ASGD)
 
+    def test_activations_factory(self):
+        class DummyModel(nn.Module):
+            @resolve_param("activation_in_head", ActivationsTypeFactory())
+            def __init__(self, activation_in_head):
+                super().__init__()
+                self.activation_in_head = activation_in_head()
+
+        model = DummyModel(activation_in_head="leaky_relu")
+        self.assertIsInstance(model.activation_in_head, nn.LeakyReLU)
+
 
 if __name__ == '__main__':
     unittest.main()