Search space auto generation

enot.autogeneration package contains functional for automatic search space generation from user pre-trained models.

Currently, auto generation supports MobileNet-like, ResNet-like YoloV5 and EfficientNet-like blocks. If your model consists of basic building blocks from these models - then you can also use this module to automate your search space construction. Trained weights from the user model are used to initialize weights of generated search space.

search_variants_from_model

generate_pruned_search_variants_model(model, search_variant_descriptors=(TransformationParameters(width_mult=1.0, kernel_size=None, depth_mult=1.0), TransformationParameters(width_mult=0.75, kernel_size=None, depth_mult=1.0), TransformationParameters(width_mult=0.5, kernel_size=None, depth_mult=1.0), TransformationParameters(width_mult=0.25, kernel_size=None, depth_mult=1.0), TransformationParameters(width_mult=0.0, kernel_size=None, depth_mult=1.0)), excluded_modules=None, concrete_args=None, keep_debug_names=False)

Generates model with search variants from source model.

Weights of source model are used for initialization of search variants weights. By default, generates model containing search variants with width multipliers (1.0, 0.75, 0.5, 0.25, 0.0).

Parameters:

  • model (torch.nn.Module) – Source PyTorch model.

  • search_variant_descriptors (tuple with TransformationParameters, optional) – Transformations which define which operations will be used as search variants. See more in Autogeneration transformation parameters. By default, uses width multipliers 1.0, 0.75, 0.5, 0.25, 0.0, and preserves kernel sizes. For Yolov5-like models choosing kernel size is not implemented.

  • excluded_modules (list with types of modules or instances of torch.nn.Module, or None, optional) – Types of modules or module instances that should be excluded from pattern matching. Default value is None, which does not exclude any module.

  • concrete_args (dict with str keys and any values, optional) – Keys in dict are names of forward parameters which we want to bind with concrete values to arguments during symbolic tracing. This is useful in cases where there are additional arguments in the forward method of model(ex. verbose, profile) or we explicitly want to commit the input flow for tracing.

  • keep_debug_names (bool, optional) – Preserve original autogenerated names. Can be useful for debugging purposes and backward compatibility. Default value is False.

Returns:

Generated model.

Return type:

torch.fx.GraphModule

Raises:

PatternsNotFoundError if autogeneration is not possible. –

Examples

>>> import torch.nn as nn
>>> from torchvision.models.mobilenetv2 import InvertedResidual
>>> from enot.models import SearchSpaceModel
>>> from enot.autogeneration import generate_pruned_search_variants_model
>>> from enot.autogeneration import TransformationParameters
>>> my_model = nn.Sequential(
...    InvertedResidual(inp=3, oup=32, stride=2, expand_ratio=6),
...    InvertedResidual(inp=32, oup=32, stride=1, expand_ratio=6),
...)
>>> generated_model = generate_pruned_search_variants_model(
...    my_model,
...    search_variant_descriptors=(
...        TransformationParameters(width_mult=0.5, kernel_size=5),
...        TransformationParameters(kernel_size=7),
...        TransformationParameters(width_mult=0.33),
...    ),
...    excluded_modules=[my_model[0]],  # Leave the first inverted residual block unchanged.
...)
>>> search_space = SearchSpaceModel(generated_model)