dynamic

Basic dynamic module class and hparam class.

Classes

DynamicModule	Base class for dynamic modules.
DynamicSpace	A class to represent all dynamic model choices over a model with multiple submodules.

class DynamicModule

Bases: Module

Base class for dynamic modules.

Dynamic modules are usually extended from torch.nn.Module’s. They dynamically support a family of torch.nn.Module’s with different architectural parameters, such as input/output channel numbers.

Dynamic modules can also be used to construct the basic searchable unit in a search space with the option to select and sample a candidate unit.

Candidate units are usually described using Hparam objects and dynamic attributes. Each hparam describes a basic searchable unit (e.g. number of output channels in DynamicConv2d). Dynamic attributes are callbacks to dynamically construct the attribute depending on the current value of the hparam(s), e.g., the weight tensor in DynamicConv2d that depends on the hparam out_channels.

In addition, dynamic modules also support registering general attributes that are removed upon export the module. This is useful for storing temporary attributes that are not hparams or existing attributes that are converted to dynamic attributes.

For a DynamicModule class that contains other dynamic modules, the class implementation should ensure only to expose hparams in the outermost class and handle other hparams internally including hparams of child modules that are exposed on their own usually (e.g. block module implementations containing DynamicLinear).

__init__(*args, **kwargs)

Initializing a dynamic module is not allowed!

classmethod convert(module)

Converts a module in-place into its dynamic counterpart by patching its class.

Parameters:

module (Module) – The module to be converted into a dynamic module.

Returns:

The converted dynamic module.

Return type:

DynamicModule

This should generally be a final method and child classes should inherit _setup() instead to customize the conversion process.

Patching is achieved by updating the __class__ attribute of the module to its dynamic counterpart. The dynamic counterpart is a subclass of the original class, hence, we ensure the module is fully compatible with the original class. Simultaneously, we can inject the corresponding dynamic behavior in a standardized and rigoruos fashion.

export()

Export self (a dynamic module) in-place and return the exported module.

The export process will remove the top-level dynamic module and replace it with the original module class. Note that the original class may be either another type of dynamic module or a vanilla nn.Module. Consequently, any methods (including properties) that are implemented in the child class will be removed. Hparams as well as dynamic and temporary attributes are handled in a special fashion, see below.

In order to ensure that the exported module is still consistent there a several mechanisms in place to handle hparams, dynamic attributes, and temporary attributes:

• Hparams of the current type are replaced with their currently active value.

  Note that we do not need to explicitly handle hparams of the parent class as they are mutually-exlusive, i.e., hparams are unique across all levels of inheritance.

• Dynamic Attributes are handled depending on whether they exist in a parent class:

  1. The same dynamic attribute exists in a parent class. In this case, the callback is folded into (“appended to”) the callback for the same dyanmic attribute of the parent class. This way we ensure that the final value of the attribute remains consistent.

  2. The dynamic attribute does not exist in a parent class. In this case, the attribute is not dynamic anymore as there are no more callbacks that could affect the value. Therefore, we simply overwrite the underlying original object with the current value and revert it to a regular attribute.

• Temporary Attributes are kept until the final export, i.e., until the resulting class is not a dynamic module anymore. This is to ensure that folded callbacks that may need access to these attributes can still access them.

Return type:

Module

extra_repr()

Generate extra_repr making sure all dynamic keys exist in self.__dict__.

Pytorch heavily uses self.__dict__ to generate extra_repr. However, we remove certain attributes from self.__dict__ so we can manage them dynamically. Temporarily, adding them back in here and removing them again afterwards.

force_assign()

Force re-assign all dynamic attributes to their current values.

Warning

Note that this method overwrittes the actual buffers and parameters! Only use in specific circumstances!!

freeze()

Restrict the hparams of the dynamic module to the orginal choices.

This is useful to enforce the behavior of the parent class.

Note

After this call, the module’s hparams can no longer be modified although the underlying type is still a dynamic module.

get_hparam(target)

Look up and return hparam (like “torch.nn.Module.get_parameter()” but for hparam).

Parameters:

target (str) –

Return type:

Hparam

modify(*args, **kwargs)

Modify the module’s dynamic choices in a standardized & scalable fashion.

This method can be overriden by the child class! While users can also directly modify the choices of individual hparams, this method should provide a way to modify a batch of dynamic modules with the same arguments, e.g., out_features_ratio for DynamicLinear.

Note that arguments of the modify method that are exposed to the user via the rule system should be specified as keyword-only arguments. When they are exposed as keyword-only arguments, the _DMRegistryCls can automatically generate the corresponding config class on the fly that lets user provide configs and then they are automatically validated before being passed to the modify method.

If possible, modify()’s keyword arguments should have default values that leave the hparams intact if not provided, e.g., one might call some_dynamic_module.modify() without any arguments and the module will remain unchanged.

named_hparams(configurable=None)

Return an iterator over all hparams of the module.

Parameters:

configurable (bool | None) – Whether to include configurable hparams.

Yields:

(name, Hparam) – Tuple containing the name and hparam.

Return type:

Iterator[Tuple[str, Hparam]]

Default behavior is to iterate over configurable and non-configurable hparams. Set configurable accordingly to only iterate over either. If configurable is set to True, only configurable symbols are iterated over. If configurable is set to False, configurable symbols are skipped over (only non-configurable symbols).

property original_cls: Type[Module]

Return the original class of the dynamic module.

reset_dynamic_attributes()

Context manager to temporarily remove any dynamic attributes and re-register values.

This context manager is intended to be used when we want to access a dynamic attribute in its original unmodified version, i.e., without this class interfering with its original value and its corresponding getattr/setattr/delattr behavior.

Upon exiting the context manager, the dynamic attributes are re-registered and the same callbacks are re-registered together with the new value.

class DynamicSpace

Bases: object

A class to represent all dynamic model choices over a model with multiple submodules.

__init__(model)

Initialize the dynamic space from the model.

Parameters:

model (Module) –

Return type:

None

config(configurable=None)

Return the config dict of all hyperparameters.

Parameters:

• model – A model that contains DynamicModule(s).

• configurable (bool | None) – None -> all hps, True -> configurable hps, False -> non-configurable hps

Returns:

A dict of (parameter_name, choice) that specifies an active subnet.

Return type:

Dict[str, Any]

convert_to_dynamic(rules, dm_registry)

Convert the model to dynamic modules according to the rules and provided registry.

Parameters:

• rules (Dict[str, Dict[str, Any] | None | Dict[str, Dict[str, Any] | None]] | None) – A dictionary containing rules for the dynamic modules.

• dm_registry (_DMRegistryCls) – A registry containing the dynamic modules to be converted to.

Returns:

A dictionary containing the converted modules with submodule names as keys and the converted dynamic modules as values.

Return type:

Dict[str, Module]

export(dm_registry)

Recursively export the module including self and return the result.

Parameters:

dm_registry (_DMRegistryCls) – A dynamic module registry to check for dynamic modules that should be exported.

Returns:

The model after exporting the dynamic modules found in the registry.

Return type:

Module

get_hparam(name)

Get the hparam with the given name.

Parameters:

name (str) –

Return type:

Hparam

is_configurable()

Check if the model has any configurable hyperparameters.

Parameters:

model – A model to be checked for DynamicModule(s) with configurable hyperparameters.

Returns:

True if the model contains DynamicModule(s) with configurable hyperparameters w/ more than one choice.

Return type:

bool

is_dynamic()

Check if any module is dynamic.

Returns:

True if the model contains DynamicModule(s).

Return type:

bool

named_dynamic_modules()

Recursively yield the name and instance of all DynamicModules.

Yields:

(name, DynamicModule) – Tuple containing the name and module.

Return type:

Iterator[Tuple[str, DynamicModule]]

named_hparams(configurable=None)

Recursively yield the name and instance of all hparams.

Parameters:

configurable (bool | None) – Whether to include configurable hparams.

Yields:

(name, Hparam) – Tuple containing the name and hparam.

Return type:

Iterator[Tuple[str, Hparam]]

Default behavior is to iterate over all hparams. If configurable is set to True, only configurable, non-duplicate symbols are iterated over.

select(config, strict=True)

Select the subnet provided by config.

If strict is set, then config must contain the exact set of keys representing both the configurable and non-configurable hparams.

Parameters:

• config (Dict[str, Any]) –

• strict (bool) –

Return type:

None

size()

Get the search space size of the model.

Returns:

A int representing the search space size of the model.

Return type:

int