search_space

An extension of the DynamicSpace to handle searchable units and related utilities.

Classes

SearchSpace

Stub for search space interface.

Functions

generate_search_space

Patch model with dynamic units and generate/return the search space.

class SearchSpace

Bases: DynamicSpace

Stub for search space interface.

__init__(model, sym_map=None)

Initialize the search space from the model and its associated symbolic map.

Parameters:

model (Module) –
sym_map (SymMap | None) –

Return type:

None

export(dm_registry=<modelopt.torch.opt.dynamic._DMRegistryCls object>)

Export with default registry.

Parameters:: dm_registry (_DMRegistryCls) –
Return type:: Module

generate(rules=None)

Generate the search space from the model + symbolic map.

Parameters:: rules (Dict[str, Dict[str, Any] | None | Dict[str, Dict[str, Any] | None]] | None) –
Return type:: None

print_summary(skipped_hparams=['kernel_size'])

Print a summary of the search space.

Parameters:: skipped_hparams (List[str]) –
Return type:: None

sample(sample_func=<function choice>)

Sample configurable hparams using the provided sample_func and return resulting config.

Parameters:

model – A configurable model that contains one or more DynamicModule(s).
sample_func (Callable[[Sequence[T]], T] | Dict[str, Callable[[Sequence[T]], T]]) –
A sampling function for hyperparameters. If a dict is provided, the key indicates a Unix shell-style wildcard pattern which is used to match the hyperparameter name. The default is random sampling corresponding to {"*": random.choice}.

Note

Later matches take precedence over earlier ones when a dict with wilcards is provided.

Returns:

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

Return type:

Dict[str, Any]

sort_parameters(hps_to_sort=None)

A graph propagation based parameter sorting algorithm.

Parameters:: hps_to_sort (Set[str] | None) – A set of hparam names to sort. If not provided or empty, all hparams will be sorted.
Return type:: None

generate_search_space(model, rules=None)

Patch model with dynamic units and generate/return the search space.

Parameters:

model (Module) – The model to patched for which we want to generate the search space.

rules (Dict[str, Dict[str, Any] | None | Dict[str, Dict[str, Any] | None]] | None) –

An optional dict specifying the rules used during the search space generation. If not provided, the largest possible search space is generated. An example rule:

rules = {
    "nn.Sequential": {"min_depth": 2},
    "nn.Conv2d": {
        "*": {
            "channels_ratio": [1],
            "channel_divisor": 16,
            "kernel_size": [],
        },
        "backbone.stages.[1-5]*.spatial_conv*": {
            "channels_ratio": [0.334, 0.5, 0.667, 1],
        },
        "det_blocks.*.spatial_conv*": {
            "channels_ratio": [0.667, 1],
        },
    },
    "nn.Linear": {},  # uses default arguments for Linear
    "nn.BatchNorm2d": {},
    "nn.SyncBatchNorm": {},
}

Note

When rules are provided, the search space generation is restricted to the rules that are provided. Other module types will not be included in the search space.

Note

Rules with None, e.g., {"nn.Sequential": None}, will also cause the conversion to be skipped just like when no rules are provided for the layer type.

Note

Instead of specifying the module type as a string, the module type can also be specified as a key directly, e.g., {nn.Sequential: None}.

Returns:

The patched model with dynamic units that describe the resulting search space.

Return type:

SearchSpace

Generally, the search space generation follows the following outline of steps:

Analyze the model to identify the dynamic units that can be patched via modelopt.torch.trace.

Patch the model with the dynamic units accordingly.

Generate a consistent search space from the patched model by ensuring consistency between dynamic units according to the symbolic map obtained from the model analysis.

Upon return of the model, the search space can be analyzed and modified via the utilities provided in modelopt.torch.dynamic.SearchSpace.