Overriding values of type Type

[DirkKuhn]

Hi,

I am using hydra_zen to configure a machine learning project. There I want to pass the type of the activation function to the Model, so that the model can instantiate it. Now to the problem: I am unable to overrride the activation function from the launch function.
Here is a minimal example:

from typing import Type
from torch import nn

from hydra_zen import to_yaml, store, builds, zen, launch
from omegaconf import DictConfig, OmegaConf

class Model:
    def __init__(self, activation_fn: Type[nn.Module]):
        self.activation_fn = activation_fn

def app(zen_cfg: DictConfig, model: Model) -> None:
    OmegaConf.resolve(zen_cfg)
    print(to_yaml(zen_cfg))

store(Model, activation_fn=nn.PReLU, group='model')

Config = builds(
    app,
    hydra_defaults=[
        {'model': 'Model'},
        '_self_'
    ]
)

if __name__ == '__main__':
    store.add_to_hydra_store()
    launch(Config, zen(app), version_base='1.3')

This does exactly what I want. The resulting zen_cfg ist

model:
  _target_: __main__.Model
  activation_fn:
    path: torch.nn.modules.activation.PReLU
    _target_: hydra_zen.funcs.get_obj
_target_: __main__.app

i.e. activation function will actually be a type.
Now, if I try to override the activation function by

launch(Config, zen(app), version_base='1.3', overrides=['model.activation_fn=nn.SELU'])

I get the zen_cfg

model:
  _target_: __main__.Model
  activation_fn: nn.SELU
_target_: __main__.app

i.e. the activation function is a string.

Is there a way to do this? Or might one of the following alternatives work: Changing the Config object directly. Overriding model.activation_fn.path and model.activation_fn._target_ directly. Creating a config for each activation function and storing them in a config group, then just overriding the used config group.

Thanks a lot!

[DirkKuhn]

I did a bit of experimenting and the approach with config groups seems to work:

from typing import Type
from torch import nn

from hydra_zen import to_yaml, store, builds, zen, launch, just
from omegaconf import DictConfig, OmegaConf

class Model:
    def __init__(self, activation_fn: Type[nn.Module]):
        self.activation_fn = activation_fn

def app(zen_cfg: DictConfig, model: Model) -> None:
    OmegaConf.resolve(zen_cfg)
    print(to_yaml(zen_cfg))

store(Model, group='model')

activation_fn_store = store(group='model/activation_fn')
activation_fn_store(just(nn.ReLU), name='relu')
activation_fn_store(just(nn.SELU), name='selu')

Config = builds(
    app,
    hydra_defaults=[
        {'model': 'Model'},
        {'model/activation_fn': 'relu'},
        '_self_'
    ]
)

if __name__ == '__main__':
    store.add_to_hydra_store()
    launch(Config, zen(app), version_base='1.3', overrides=['model/activation_fn=selu'])

results in

model:
  _target_: __main__.Model
  activation_fn:
    path: torch.nn.modules.activation.SELU
    _target_: hydra_zen.funcs.get_obj
_target_: __main__.app

Please tell me there is a better way to do this! Or should I approach the problem completely differently?
Again thanks so much!

[rsokl]

Yep, config groups is the right way to solve this issue! Nice work!

(Sorry for the slow responses - this is a very busy week for me)

[DirkKuhn]

No worries, I will probably spend a lot of time here anyway.
Thank you for creating and publishing hydra_zen :)