# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Can't have annotations on because PyTorch JIT doesn't support it.
# from __future__ import annotations
__all__ = [
'recursive_freeze', 'MutableModule', 'ModelSpace', 'ParametrizedModule'
]
import copy
import itertools
import inspect
import logging
from typing import Any, Callable, Type, Iterable, Dict, Tuple, List, Optional, TypeVar
import torch
from torch import nn
from nni.mutable import (
Mutable, LabeledMutable, Sample, SampleValidationError,
ensure_frozen, frozen_context, label_scope, frozen
)
from nni.nas.space import BaseModelSpace, current_model, model_context
from nni.common.serializer import SerializableObject, _formulate_arguments, _copy_class_wrapper_attributes
_logger = logging.getLogger(__name__)
T = TypeVar('T', bound=Mutable)
def recursive_freeze(module: nn.Module, sample: Dict[str, Any], include_root: bool = True) -> Tuple[nn.Module, bool]:
"""Recursively freeze all the mutables in a module.
The sample is a dictionary of parameters.
If ``include_root`` is true, it will call ``module.freeze(sample)`` if module is a mutable by itself.
Returns
-------
A tuple of (frozen_module, replaced).
If replaced is true, the child module has been replaced.
"""
if include_root and isinstance(module, Mutable):
return module.freeze(sample), True
# Examine all the children.
replaced_dict = {}
for name, child_module in module.named_children():
if isinstance(child_module, Mutable):
new_module, replaced = child_module.freeze(sample), True
else:
new_module, replaced = recursive_freeze(child_module, sample)
if replaced:
replaced_dict[name] = new_module
# If some children is replaced.
if replaced_dict:
module = copy.deepcopy(module)
for name, new_module in replaced_dict.items():
setattr(module, name, new_module)
return module, True
return module, False
[文档]
class MutableModule(Mutable, nn.Module):
"""
PyTorch module, but with uncertainties.
This base class provides useful tools to handle search spaces built on top of PyTorch modules,
including methods like :meth:`simplify`, :meth:`freeze`.
:class:`MutableModule` can have dangling mutables registered on it via :meth:`add_mutable`.
"""
def __new__(cls, *args, **kwargs):
# The purpose of __new__ is to intercept module creation,
# and create other types of modules when necessary.
arch = current_model()
if cls.should_invoke_fixed_module() and arch is not None:
# If within a fixed_arch context, create the frozen module.
# It must return a object with different type, or else infinite recursion will happen.
return cls.create_fixed_module(arch, *args, **kwargs) # type: ignore
else:
return super().__new__(cls)
[文档]
@classmethod
def should_invoke_fixed_module(cls):
"""Call ``create_fixed_module()`` when fixed-arch context is detected.
Typically this should be enabled. Otherwise the arch context might not be correctly handled.
In cases where this flag is disabled, remember to detect arch context and
manually freeze things in ``__init__``, or confirm that it's a composite module
and nothing needs to be frozen.
By default, it returns true when :meth:`create_fixed_module` is overridden.
"""
return not getattr(cls.create_fixed_module, '_notimplemented', False)
[文档]
def add_mutable(self, mutable: T) -> T:
"""Register a mutable to this module.
This is often used to add dangling variables that are not parameters of any :class:`ParametrizedModule`.
If the mutable is also happens to be a submodule of type :class:`MutableModule`,
it can be registered in the same way as PyTorch (i.e., ``self.xxx = mutable``). No need to add it again here.
Examples
--------
In practice, this method is often used together with :func:`~nni.mutable.ensure_frozen`.
>>> class MyModule(MutableModule):
... def __init__(self):
... super().__init__()
... token_size = nni.choice('t', [4, 8, 16]) # Categorical variable here
... self.add_mutable(token_size) # Register the mutable to this module.
... real_token_size = ensure_frozen(token_size) # Real number. 4 during dry run. 4, 8 or 16 during search.
... self.token = nn.Parameter(torch.randn(real_token_size, 1))
.. tip::
Note that :func:`~nni.mutable.ensure_frozen` must be used under a :func:`~nni.mutable.frozen_context`.
The easiest way to do so is to invoke it within initialization of a :class:`ModelSpace`.
Warnings
--------
Arbitrary :meth:`add_mutable` is not supported for :class:`~nni.nas.space.GraphModelSpace`.
"""
# NOTE:
# This method can be potentially useful for every mutable. But I leave it here for the following reasons:
# 1. If `_mutables` is to be put into base class, all subclasses need to call `super().__init__()`.
# 2. Most of other mutables are "leaves". They do not have children. I don't agree that all mutables are "containers".
# 3. To support nested mutables (like NestedCategorical), designs are needed (e.g., can constants be added via `add_mutable`?).
# 4. The interface can be easily extended to base class if needed in future, without breaking backward compatibility.
if isinstance(mutable, MutableModule):
# If the mutable is also a module, it will be added automatically.
return mutable
if not isinstance(mutable, Mutable):
raise TypeError(f'Expected Mutable, got {type(mutable)}: {mutable}')
# Dynamically create it here to avoid an __init__ call.
# Also avoiding dependencies for __init__ in ParametrizedModule.
if not hasattr(self, '_mutables'):
self._mutables: List[Mutable] = []
self._mutables.append(mutable)
# Disable dry run when:
# 1. Inside a model_context; the `current_model()` needs to be read-only.
# 2. ensure_frozen_strict is disabled. The `default()` can be called when `ensure_frozen()` is called.
if current_model() is None:
# We will "dry run" the mutable here, to obtain a default value for it.
context = frozen_context.current()
if context is None:
if frozen._ENSURE_FROZEN_STRICT:
_logger.warning(
'No context found when `add_mutable()`. '
'You are probably adding a MutableModule outside the `__init__` of a ModelSpace. '
'This can possibly make the MutableModule untrackable and inconsistent: %s',
mutable
)
else:
context_before_keys = set(context.keys())
try:
mutable.robust_default(context)
except ValueError:
_logger.error('Error when trying to dry run the mutable: %r. It could be conflicted with current context: %s.',
mutable, context)
raise
frozen_context.update(
{key: value for key, value in context.items() if key not in context_before_keys}
)
return mutable
@torch.jit.unused # Must be marked as unused here. Otherwise JIT will look into Mutable and fail.
@property
def mutables(self) -> List[Mutable]:
"""Mutables that are dangling under this module.
Normally this is all the mutables that are registered via :meth:`MutableModule.add_mutable`.
"""
if not hasattr(self, '_mutables'):
self._mutables: List[Mutable] = []
return self._mutables
# This is actually a classmethod, but decorated afterwards to assign `_notimplemented` attribute.
# @classmethod
[文档]
def create_fixed_module(cls, sample: dict, *args, **kwargs) -> nn.Module: # type: ignore
"""
The classmethod is to create a brand new module with fixed architecture.
The parameter ``sample`` is a dict with the exactly same format as ``sample`` in :meth:`freeze`.
The difference is that when :meth:`create_fixed_module` is called,
there is no :class:`MutableModule` instance created yet.
Thus it can be useful to simplify the creation of a fixed module,
by saving the cost of creating a :class:`MutableModule` instance and immediately :meth:`freeze` it.
If automatic label generation (e.g., :func:`~nni.mutable.auto_label`) is used in ``__init__``,
the same number of labels should be generated in this method.
Otherwise it will mess up the global label counter, and potentially affect the label of successive modules.
By default, this method has a not-implemented flag, and :meth:`should_invoke_fixed_module` will return ``False``
based on this flag.
"""
raise NotImplementedError('create_fixed_module() must be implemented when `custom_fixed_module_creation` is set to true.')
create_fixed_module._notimplemented = True
create_fixed_module = classmethod(create_fixed_module) # type: ignore
def check_contains(self, sample: Sample) -> Optional[SampleValidationError]:
for mutable in self.mutables:
exception = mutable.check_contains(sample)
if exception is not None:
return exception
for name, module in self.named_mutable_descendants():
exception = module.check_contains(sample)
if exception is not None:
exception.paths.append(name)
return exception
return None
def leaf_mutables(self, is_leaf: Callable[[Mutable], bool]) -> Iterable[LabeledMutable]:
for mutable in self.mutables:
yield from mutable.leaf_mutables(is_leaf)
for module in self.mutable_descendants():
yield from module.leaf_mutables(is_leaf)
[文档]
def mutable_descendants(self) -> Iterable['MutableModule']:
""":meth:`named_mutable_descendants` without names."""
for _, module in self.named_mutable_descendants():
yield module
[文档]
def named_mutable_descendants(self) -> Iterable[Tuple[str, 'MutableModule']]:
"""Traverse the module subtree, find all descendants that are :class:`MutableModule`.
- If a child module is :class:`MutableModule`, return it directly, and its subtree will be ignored.
- If not, it will be recursively expanded, until :class:`MutableModule` is found.
"""
def _iter(name: str, module: nn.Module) -> Iterable[Tuple[str, MutableModule]]:
for subname, child in module.named_children():
name_ = name + '.' + subname if name else subname
if isinstance(child, MutableModule):
yield name_, child
else:
yield from _iter(name_, child)
yield from _iter('', self)
[文档]
def freeze(self, sample: Dict[str, Any]) -> nn.Module:
"""Return a frozen version of current mutable module.
Some sub-modules can be possibly deep-copied.
If mutables are added to the module via :meth:`add_mutable`, this method must be implemented.
Otherwise, it will simply look at the children modules and freeze them recursively.
:meth:`freeze` of subclass is encouraged to keep the original weights at best effort,
but no guarantee is made, unless otherwise specified.
"""
if self.mutables:
raise NotImplementedError(
'freeze() must be implemented when mutables have been registered to the module. '
'A simple way to implement this is to recreate the current module after set the frozen context:\n\n'
' with nni.nas.space.model_context(sample):\n'
' return self.__class__(*args, **kwargs)\n\n'
'Here, `args` and `kwargs` are the arguments to the constructor of the module.'
)
return recursive_freeze(self, sample, include_root=False)[0]
def __repr__(self):
return nn.Module.__repr__(self)
class TraceableMixin(Mutable):
"""
This is actually another implementation of ``nni.trace()``.
In this implementation, no decorator is needed. Users only need to inherit this class.
``is_traceable()`` check will still pass.
For now, everything traced is on ``__init__``.
Classes with customized ``__new__`` are not supported.
"""
# Model space also needs to record init arguments, for recovering from another process.
trace_args: tuple
trace_kwargs: Dict[str, Any]
# Useful in getting the signature of the original class __init__.
_init_wrapped: Optional[Callable[..., None]] = None
@torch.jit.ignore # type: ignore
def save_init_arguments(self, *args, **kwargs) -> None:
self.trace_args = tuple(args)
self.trace_kwargs = dict(kwargs)
@torch.jit.ignore # type: ignore
def auto_save_init_arguments(self, *args, **kwargs) -> None:
"""Save init arguments into ``trace_args`` and ``trace_kwargs``.
Skip when ``trace_args`` and ``trace_kwargs`` are already set,
which could be possibly due to subclassing / inheritance.
"""
# Stop saving when it's already saved, possibly in subclass's __init__.
if not hasattr(self, 'trace_args') and not hasattr(self, 'trace_kwargs'):
# If both super-class and sub-class calls this, subclass should first call this.
init_fn = getattr(self.__class__, '_init_wrapped', self.__class__.__init__)
args_, kwargs_ = _formulate_arguments(init_fn, args, kwargs, True, True)
self.save_init_arguments(*args_, **kwargs_)
@torch.jit.unused
@property
def trace_symbol(self):
return self.__class__
@torch.jit.unused
@property
def args(self) -> Dict[str, Any]:
"""All arguments that are used to construct the instance,
including arguments passed to ``__init__``, as well as arguments with default value.
Positional arguments are not supported.
"""
if self.trace_args:
raise RuntimeError('args is not available when positional argument is not empty.')
rv = dict(self.trace_kwargs)
# Add the arguments that are in signature and with default value.
init_fn = getattr(self.__class__, '_init_wrapped', self.__class__.__init__)
for param in inspect.signature(init_fn).parameters.values():
if param.default is not param.empty and param.name not in rv:
rv[param.name] = param.default
return rv
@torch.jit.ignore # type: ignore
def trace_copy(self):
"""Returns a different object here. All the model-specific details will be thrown away."""
return SerializableObject(self.__class__, list(self.trace_args), self.trace_kwargs)
[文档]
class ModelSpace(
TraceableMixin,
MutableModule,
BaseModelSpace,
):
"""
The base class for model search space based on PyTorch.
The out-est module should inherit this class.
Model space is written as PyTorch module for the convenience of writing code.
It's not a real PyTorch model, and shouldn't be used as one for most cases.
Most likely, the forward of :class:`ModelSpace` is a dry run of an arbitrary model in the model space.
But since there is no guarantee on which model will be chosen, and the behavior is not well tested,
it's only used for sanity check and tracing the space, and its semantics are not well-defined.
Similarly for ``state_dict`` and ``load_state_dict``.
Users should bear in mind that :class:`ModelSpace` is NOT a one-shot supernet,
directly exporting its weights are unreliable and prone to error.
Use :ref:`one-shot strategies <one-shot-nas>` to mutate the model space into a supernet for such needs.
Mutables in model space **must all be labeled manually**, unless a label prefix is provided.
Every model space can have a label prefix, which is used to provide a stable automatic label generation.
For example, if the label prefix is ``model``, all the mutables initialized in a subclass of ModelSpace
(in ``__init__`` function of itself and submodules, to be specific),
will be automatically labeled with a prefix ``model/``.
The label prefix can be manually specified upon definition of the class::
class MyModelSpace(ModelSpace, label_prefix='backbone'):
def __init__(self):
super().__init__()
self.choice = self.add_mutable(nni.choice('depth', [2, 3, 4]))
print(self.choice.label) # backbone/choice
Notes
-----
The ``__init__`` implementation of :class:`ModelSpace` is in :func:`model_space_init_wrapper`.
"""
# JIT can't parse label_scope. Deliberately ignore it to make JIT happy.
# _label_scope: label_scope
_label_prefix: Optional[str]
"""The label prefix of the model space."""
def __init_subclass__(cls, disable_init_wrapper: bool = False, label_prefix: Optional[str] = None, **kwargs) -> None:
# The init wrapper can be turned off in tricky cases.
if not disable_init_wrapper:
cls._init_wrapped = cls.__init__
cls.__init__ = model_space_init_wrapper(cls.__init__)
cls._label_prefix = label_prefix
[文档]
@classmethod
def load_searched_model(cls, name: str, pretrained: bool = False, download: bool = False, progress: bool = True) -> nn.Module:
"""Load a pre-searched model with given name."""
raise NotImplementedError('`load_searched_model` is not implemented, which means that no pre-searched model is available.')
class strict_label_scope(label_scope):
"""A strict label scope that raises error when label is not manually specified."""
def next_label(self) -> str:
raise ValueError('Label must be specified manually in NAS, or provide a `label_prefix` to the model space.')
@property
def path(self) -> Optional[List[str]]:
"""A strict label scope is only used for label checking. It shouldn't have its own name."""
if self._path is None:
return self._path
return self._path[:-1]
def model_space_init_wrapper(original_init_fn: Callable[..., None]) -> Callable[..., None]:
"""Wrap the ``__init__`` inside a model namespace.
This should only be wrapped on *subclasses* of :class:`ModelSpace`.
It could possibly be wrapped multiple times when subclass of :class:`ModelSpace` is further subclassed.
Currently, nested label namespace and dry run context will be created.
"""
def init_with_context(self: ModelSpace, *args, **kwargs):
arch = current_model()
if arch is None:
# If no arch is set, we are in dry run mode.
# This is the only place where we set the dry run mode.
# The context will record the sample generated during ``__init__``.
# Some mutables use :func:`ensure_frozen` to freeze the sample.
# The context is to make sure same label always return the same choice.
self._frozen_context = frozen_context()
with self._frozen_context:
return original_init_fn(self, *args, **kwargs)
else:
# Already in context.
self._frozen_context = frozen_context.top_context()
return original_init_fn(self, *args, **kwargs)
def new_init(self: ModelSpace, *args, **kwargs) -> None:
# Save init arguments first. Skips if already saved.
self.auto_save_init_arguments(*args, **kwargs)
if not hasattr(self, '_label_scope'):
if self._label_prefix is not None:
self._label_scope = label_scope(self._label_prefix)
else:
self._label_scope = strict_label_scope('_unused_') # the name is not used
if hasattr(self, '_label_scope') and not self._label_scope.activated: # type: ignore
# Has a label scope but it's not activated. Create a "with".
with self._label_scope: # type: ignore
return init_with_context(self, *args, **kwargs)
else:
return init_with_context(self, *args, **kwargs)
return new_init
[文档]
class ParametrizedModule(
TraceableMixin,
MutableModule,
):
"""
Subclass of :class:`MutableModule` supports mutables as initialization parameters.
One important feature of :class:`ParametrizedModule` is that
it automatically freeze the mutable arguments passed to ``__init__``.
This is for the convenience as well as compatibility with existing code::
class MyModule(ParametrizedModule):
def __init__(self, x):
super().__init__()
self.t = x # Will be a fixed number, e.g., 3.
MyModule(nni.choice('choice1', [1, 2, 3]))
Note that the mutable arguments need to be directly posed as arguments to ``__init__``.
They can't be hidden in a list or dict.
If users want to make a 3rd-party module *parametrized*,
it's recommended to do the following (taking ``nn.Conv2d`` as an example):
>>> class ParametrizedConv2d(ParametrizedModule, nn.Conv2d, wraps=nn.Conv2d):
... pass
>>> conv = ParametrizedConv2d(3, nni.choice('out', [8, 16]))
>>> conv
>>> conv.out_channels
8
>>> conv.args['out_channels']
Categorical([8, 16], label='out')
>>> conv.freeze({'out': 16})
Conv2d(3, 16, kernel_size=(1, 1), stride=(1, 1))
.. tip::
The parametrized version of modules in ``torch.nn`` are already provided in ``nni.nas.nn.pytorch``.
Every class is prefixed with ``Mutable``.
For example, :class:`nni.nas.nn.pytorch.MutableConv2d`` is a parametrized version of ``torch.nn.Conv2d``.
Attributes
----------
args
The arguments used to initialize the module.
Since :class:`ParametrizedModule` will hijack the init arguments before passing to ``__init__``,
this is the only recommended way to retrieve the original init arguments back.
Warnings
--------
:class:`ParametrizedModule` can be nested.
It's also possible to put arbitrary mutable modules inside a :class:`ParametrizedModule`.
But be careful if the inner mutable modules are dependant on the parameters of :class:`ParametrizedModule`,
because NNI can't handle cases where the mutables are a dynamically changing after initialization.
For example, the following snippet is WRONG::
class MyModule(ParametrizedModule):
def __init__(self, x):
if x == 0:
self.mutable = self.add_mutable(nni.choice('a', [1, 2, 3]))
else:
self.mutable = self.add_mutable(nni.choice('b', [4, 5, 6]))
module = MyModule(nni.choice('x', [0, 1]))
"""
_nni_basic_unit: bool = True
"""This is only used in parsing graph. When ``_nni_basic_unit`` is true,
the graph converter will consider the module as a primitive,
and stop digging into the module to parse the inner modules.
It's considered internal and should not be used by users.
"""
_bound_type: Optional[Type] = None
"""The bounded class without :class:`ParametrizedModule`. Used to initialize a pure fixed instance.
Set to none if :class:`ParametrizedModule` is the only superclass of the defined class."""
@classmethod
def should_invoke_fixed_module(cls) -> bool:
return cls._bound_type is not None
@torch.jit.ignore # type: ignore
def __init_subclass__(
cls,
disable_init_wrapper: bool = False,
wraps: Optional[Type] = None,
copy_wrapped: bool = False,
pure_fixed_module: Optional[bool] = None,
**kwargs
) -> None:
# The init wrapper can be turned off in tricky cases.
if not disable_init_wrapper:
if wraps:
cls.__wrapped__ = wraps # type: ignore
cls._init_wrapped = wraps.__init__
else:
cls._init_wrapped = cls.__init__
cls.__init__ = parametrized_module_init_wrapper(cls.__init__)
# Copy some attributes from the wrapped class, so that the wrapped class looks exactly like the initial one.
# Useful in scenarios where module names are used to identify the class (e.g., graph space).
if copy_wrapped:
if wraps is None:
raise ValueError('`wraps` should be specified when `copy_wrapped` is set to True.')
_copy_class_wrapper_attributes(copy_wrapped, cls)
if pure_fixed_module is None:
pure_fixed_module = wraps is not None
if pure_fixed_module:
if wraps is None:
raise ValueError('`wraps` should be specified when `pure_fixed_module` is set to True.')
cls._bound_type = wraps
@classmethod
def create_fixed_module(cls, sample, *args, **kwargs) -> Any:
assert cls._bound_type is not None, 'Cannot create fixed module for a class that is not bound to a fixed type.'
args, kwargs = cls.freeze_init_arguments(sample, *args, **kwargs)
with model_context(sample): # A context should already exists. But it doesn't harm to create a new one.
return cls._bound_type(*args, **kwargs) # type: ignore # pylint: disable=not-callable
[文档]
def freeze(self, sample: Dict[str, Any]) -> nn.Module:
"""Freeze all the mutable arguments in init.
Note that a brand new module will be created, and all previous weights will be lost.
Supernet must be created with one-shot strategies if you want to keep the weights.
"""
args, kwargs = self.freeze_init_arguments(sample, *self.trace_args, **self.trace_kwargs)
with model_context(sample): # provide a context for nested mutable modules
if self._bound_type is not None:
return self._bound_type(*args, **kwargs) # type: ignore # pylint: disable=not-callable
else:
return self.__class__(*args, **kwargs)
[文档]
@staticmethod
def freeze_init_arguments(sample: Optional[Sample], *args, **kwargs) -> Tuple[tuple, dict]:
"""Freeze the init arguments with the given context, and return the frozen arguments."""
args_ = tuple(ensure_frozen(arg, sample=sample) for arg in args)
kwargs_ = {kw: ensure_frozen(arg, sample=sample) for kw, arg in kwargs.items()}
return args_, kwargs_
def __repr__(self):
params = []
for value in self.trace_args:
params.append(repr(value))
for name, value in self.trace_kwargs.items():
params.append(f'{name}={repr(value)}')
return f'{self.__class__.__name__}({", ".join(params)})'
def parametrized_module_init_wrapper(original_init_fn: Callable[..., None]) -> Callable[..., None]:
"""Wrap the ``__init__`` for :class:`ParametrizedModule`.
It returns a function, which first records all the parameters passed to ``__init__``,
then ensures all the parameters are frozen, and finally calls the original ``__init__``.
"""
def new_init(self: ParametrizedModule, *args, **kwargs):
arch = current_model()
if arch is not None:
# This is the case where fixed-arch context exists,
# yet ParametrizedModule doesn't have a bound type.
args, kwargs = self.freeze_init_arguments(arch, *args, **kwargs)
return original_init_fn(self, *args, **kwargs)
# This is the case where fixed-arch context doesn't exist.
self.auto_save_init_arguments(*args, **kwargs)
for arg in itertools.chain(args, kwargs.values()):
if isinstance(arg, Mutable):
self.add_mutable(arg)
else:
_warn_if_nested_mutable(arg, self.__class__.__name__)
# Sometimes, arguments will be hijacked to make the inner wrapped class happy.
# For example Conv2d(choice([3, 5, 7])) should be Conv2d(3) instead,
# because Conv2d doesn't recognize choice([3, 5, 7]).
args, kwargs = self.freeze_init_arguments(None, *args, **kwargs)
return original_init_fn(self, *args, **kwargs)
return new_init
def _warn_if_nested_mutable(obj: Any, cls_name: str) -> None:
# Warn for cases like MutableConv2d(kernel_size=(nni.choice([3, 5]), nni.choice([3, 5])))
# This is not designed to be reliable, but only to be user-friendly.
def _iter(o):
if isinstance(o, Mutable):
_logger.warning(f'Found a nested mutable {o} in parameter {obj} of class {cls_name}. '
'This is not recommended, because the mutable will not be tracked. '
'Please use MutableList, MutableDict instead, or write every options in a `nni.choice`.')
else:
if isinstance(o, (list, tuple, set)):
for item in o:
_iter(item)
elif isinstance(o, dict):
for value in o.values():
_iter(value)
_iter(obj)