fastnas

Module implementing fasnas pruning algorithm for search.

Classes

`BinarySearcher`	An iterative searcher that uses binary search to find the best configuration.
`FastNASConfig`	Configuration for the `"fastnas"` mode.
`FastNASModeDescriptor`	Class to describe the `"fastnas"` mode.
`FastNASPatchManager`	A patch manager for FastNAS (same as AutoNAS except no sampling during training).

Functions

`convert_fastnas_searchspace`	Convert search space for FastNAS mode with correct patch manager.
`restore_fastnas_searchspace`	Restore search space for FastNAS mode with correct patch manager.

class BinarySearcher

Bases: IterativeSearcher

An iterative searcher that uses binary search to find the best configuration.

after_step()

Update boundaries of the interval after recursing.

Return type:: None

before_search()

Build sensitivity map before search that we use to approximate the cost function.

Return type:: None

before_step()

Check what the middle value is to determine where we recurse.

Return type:: None

property default_state_dict: dict[str, Any]: We also store the sensitivity map and related arguments.

early_stop()

Early stop if the interval is small enough.

Return type:: bool

property hparam_names_for_search: set[str]: We can only optimize over certain types of hparams in binary search.

property hparam_types_for_search: tuple[type]: We can only optimize over certain types of hparams in binary search.

load_search_checkpoint()

We only want to load sensitivity map and original_score here and keep the rest.

Return type:: bool

max_degrade: float

middle_value: float

min_degrade: float

original_score: float

sample()

Check in which interval we should recurse and sets the corresponding subnet.

Return type:: dict[str, Any]

sensitivity_map: dict[str, dict[int, float]]

class FastNASConfig

Bases: ModeloptBaseRuleConfig

Configuration for the "fastnas" mode.

model_config = {'extra': 'allow', 'validate_assignment': True}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

nn_batchnorm1d: DynamicBatchNorm1dConfig | None | dict[str, DynamicBatchNorm1dConfig | None]

nn_batchnorm2d: DynamicBatchNorm2dConfig | None | dict[str, DynamicBatchNorm2dConfig | None]

nn_batchnorm3d: DynamicBatchNorm3dConfig | None | dict[str, DynamicBatchNorm3dConfig | None]

nn_conv1d: DynamicConv1dConfig | None | dict[str, DynamicConv1dConfig | None]

nn_conv2d: DynamicConv2dConfig | None | dict[str, DynamicConv2dConfig | None]

nn_conv3d: DynamicConv3dConfig | None | dict[str, DynamicConv3dConfig | None]

nn_convtranspose1d: DynamicConvTranspose1dConfig | None | dict[str, DynamicConvTranspose1dConfig | None]

nn_convtranspose2d: DynamicConvTranspose2dConfig | None | dict[str, DynamicConvTranspose2dConfig | None]

nn_convtranspose3d: DynamicConvTranspose3dConfig | None | dict[str, DynamicConvTranspose3dConfig | None]

nn_groupnorm: DynamicGroupNormConfig | None | dict[str, DynamicGroupNormConfig | None]

nn_instancenorm1d: DynamicInstanceNorm1dConfig | None | dict[str, DynamicInstanceNorm1dConfig | None]

nn_instancenorm2d: DynamicInstanceNorm2dConfig | None | dict[str, DynamicInstanceNorm2dConfig | None]

nn_instancenorm3d: DynamicInstanceNorm3dConfig | None | dict[str, DynamicInstanceNorm3dConfig | None]

nn_layernorm: DynamicLayerNormConfig | None | dict[str, DynamicLayerNormConfig | None]

nn_linear: DynamicLinearConfig | None | dict[str, DynamicLinearConfig | None]

nn_syncbatchnorm: DynamicSyncBatchNormConfig | None | dict[str, DynamicSyncBatchNormConfig | None]

class FastNASModeDescriptor

Bases: ModeDescriptor

Class to describe the "fastnas" mode.

The properties of this mode can be inspected via the source code.

property config_class: type[ModeloptBaseConfig]: Specifies the config class for the mode.

property convert: Callable[[Module, ModeloptBaseConfig], tuple[Module, dict[str, Any]]] | Callable[[Module, ModeloptBaseConfig, Any], tuple[Module, dict[str, Any]]]: The mode’s entrypoint for converting a model.

property export_mode: str | None: The mode that corresponds to the export mode of this mode.

property name: str: Returns the value (str representation) of the mode.

property next_modes: set[str] | None: Modes that must immediately follow this mode.

property restore: Callable[[Module, ModeloptBaseConfig, dict[str, Any]], Module]: The mode’s entrypoint for restoring a model.

property search_algorithm: type[BaseSearcher]: Specifies the search algorithm to use for this mode (if any).

property update_for_new_mode: Callable[[Module, ModeloptBaseConfig, dict[str, Any]], None]: The mode’s entrypoint for updating the models state before new mode.

property update_for_save: Callable[[Module, ModeloptBaseConfig, dict[str, Any]], None]: The mode’s entrypoint for updating the models state before saving.

class FastNASPatchManager

Bases: AutoNASPatchManager

A patch manager for FastNAS (same as AutoNAS except no sampling during training).

property sample_during_training: Indicates whether we should sample a new subnet during training.

convert_fastnas_searchspace(model, config)

Convert search space for FastNAS mode with correct patch manager.

Parameters:

model (Module)
config (ModeloptBaseConfig)

Return type:

tuple[Module, dict[str, Any]]

restore_fastnas_searchspace(model, config, metadata)

Restore search space for FastNAS mode with correct patch manager.

Parameters:

model (Module)
config (ModeloptBaseConfig)
metadata (dict[str, Any])

Return type:

Module