nni/algorithms/nas/pytorch/classic_nas/mutator.py [16:80]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
logger = logging.getLogger(__name__)

NNI_GEN_SEARCH_SPACE = "NNI_GEN_SEARCH_SPACE"
LAYER_CHOICE = "layer_choice"
INPUT_CHOICE = "input_choice"


def get_and_apply_next_architecture(model):
    """
    Wrapper of :class:`~nni.nas.pytorch.classic_nas.mutator.ClassicMutator` to make it more meaningful,
    similar to ``get_next_parameter`` for HPO.

    It will generate search space based on ``model``.
    If env ``NNI_GEN_SEARCH_SPACE`` exists, this is in dry run mode for
    generating search space for the experiment.
    If not, there are still two mode, one is nni experiment mode where users
    use ``nnictl`` to start an experiment. The other is standalone mode
    where users directly run the trial command, this mode chooses the first
    one(s) for each LayerChoice and InputChoice.

    Parameters
    ----------
    model : nn.Module
        User's model with search space (e.g., LayerChoice, InputChoice) embedded in it.
    """
    ClassicMutator(model)


class ClassicMutator(Mutator):
    """
    This mutator is to apply the architecture chosen from tuner.
    It implements the forward function of LayerChoice and InputChoice,
    to only activate the chosen ones.

    Parameters
    ----------
    model : nn.Module
        User's model with search space (e.g., LayerChoice, InputChoice) embedded in it.
    """

    def __init__(self, model):
        super(ClassicMutator, self).__init__(model)
        self._chosen_arch = {}
        self._search_space = self._generate_search_space()
        if NNI_GEN_SEARCH_SPACE in os.environ:
            # dry run for only generating search space
            self._dump_search_space(os.environ[NNI_GEN_SEARCH_SPACE])
            sys.exit(0)

        if trial_env_vars.NNI_PLATFORM is None:
            logger.warning("This is in standalone mode, the chosen are the first one(s).")
            self._chosen_arch = self._standalone_generate_chosen()
        else:
            # get chosen arch from tuner
            self._chosen_arch = nni.get_next_parameter()
            if self._chosen_arch is None:
                if trial_env_vars.NNI_PLATFORM == "unittest":
                    # happens if NNI_PLATFORM is intentionally set, e.g., in UT
                    logger.warning("`NNI_PLATFORM` is set but `param` is None. Falling back to standalone mode.")
                    self._chosen_arch = self._standalone_generate_chosen()
                else:
                    raise RuntimeError("Chosen architecture is None. This may be a platform error.")
        self.reset()

    def _sample_layer_choice(self, mutable, idx, value, search_space_item):
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



nni/algorithms/nas/tensorflow/classic_nas/mutator.py [18:79]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
logger = logging.getLogger(__name__)

NNI_GEN_SEARCH_SPACE = "NNI_GEN_SEARCH_SPACE"
LAYER_CHOICE = "layer_choice"
INPUT_CHOICE = "input_choice"


def get_and_apply_next_architecture(model):
    """
    Wrapper of :class:`~nni.nas.tensorflow.classic_nas.mutator.ClassicMutator` to make it more meaningful,
    similar to ``get_next_parameter`` for HPO.
    Tt will generate search space based on ``model``.
    If env ``NNI_GEN_SEARCH_SPACE`` exists, this is in dry run mode for
    generating search space for the experiment.
    If not, there are still two mode, one is nni experiment mode where users
    use ``nnictl`` to start an experiment. The other is standalone mode
    where users directly run the trial command, this mode chooses the first
    one(s) for each LayerChoice and InputChoice.
    Parameters
    ----------
    model : nn.Module
        User's model with search space (e.g., LayerChoice, InputChoice) embedded in it.
    """
    ClassicMutator(model)


class ClassicMutator(Mutator):
    """
    This mutator is to apply the architecture chosen from tuner.
    It implements the forward function of LayerChoice and InputChoice,
    to only activate the chosen ones.
    Parameters
    ----------
    model : nn.Module
        User's model with search space (e.g., LayerChoice, InputChoice) embedded in it.
    """

    def __init__(self, model):
        super(ClassicMutator, self).__init__(model)
        self._chosen_arch = {}
        self._search_space = self._generate_search_space()
        if NNI_GEN_SEARCH_SPACE in os.environ:
            # dry run for only generating search space
            self._dump_search_space(os.environ[NNI_GEN_SEARCH_SPACE])
            sys.exit(0)

        if trial_env_vars.NNI_PLATFORM is None:
            logger.warning("This is in standalone mode, the chosen are the first one(s).")
            self._chosen_arch = self._standalone_generate_chosen()
        else:
            # get chosen arch from tuner
            self._chosen_arch = nni.get_next_parameter()
            if self._chosen_arch is None:
                if trial_env_vars.NNI_PLATFORM == "unittest":
                    # happens if NNI_PLATFORM is intentionally set, e.g., in UT
                    logger.warning("`NNI_PLATFORM` is set but `param` is None. Falling back to standalone mode.")
                    self._chosen_arch = self._standalone_generate_chosen()
                else:
                    raise RuntimeError("Chosen architecture is None. This may be a platform error.")
        self.reset()

    def _sample_layer_choice(self, mutable, idx, value, search_space_item):
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -