Dubhe/dubhe-tadl/classic_nas/mutator.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import json
import logging
import os
import sys
sys.path.append('..'+ '/' + '..')
import torch



from pytorch.mutables import LayerChoice, InputChoice, MutableScope
from pytorch.mutator import Mutator
import numpy as np
import random
logger = logging.getLogger(__name__)

NNI_GEN_SEARCH_SPACE = "NNI_GEN_SEARCH_SPACE"
NNI_PLATFORM = "GPU"
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,trial_id,selected_path,search_space_path,load_selected_space=False):
        super(ClassicMutator, self).__init__(model)
        self._chosen_arch = {}
        self._search_space = self._generate_search_space()
        self.trial_id = trial_id
        #if NNI_GEN_SEARCH_SPACE in os.environ:
            # dry run for only generating search space
        self._dump_search_space(search_space_path)
            #sys.exit(0)

        if load_selected_space:
            logger.warning("load selected space.")
            self._chosen_arch = self.load_selected_space
        else:
            # get chosen arch from tuner
            self._chosen_arch = self.random_generate_chosen()


        self._generate_selected_space(selected_path)
        self.reset()

    def _sample_layer_choice(self, mutable, idx, value, search_space_item):
        """
        Convert layer choice to tensor representation.

        Parameters
        ----------
        mutable : Mutable
        idx : int
            Number `idx` of list will be selected.
        value : str
            The verbose representation of the selected value.
        search_space_item : list
            The list for corresponding search space.
        """
        # doesn't support multihot for layer choice yet
        onehot_list = [False] * len(mutable)
        assert 0 <= idx < len(mutable) and search_space_item[idx] == value, \
            "Index '{}' in search space '{}' is not '{}'".format(idx, search_space_item, value)
        onehot_list[idx] = True
        return torch.tensor(onehot_list, dtype=torch.bool)  # pylint: disable=not-callable

    def _sample_input_choice(self, mutable, idx, value, search_space_item):
        """
        Convert input choice to tensor representation.

        Parameters
        ----------
        mutable : Mutable
        idx : int
            Number `idx` of list will be selected.
        value : str
            The verbose representation of the selected value.
        search_space_item : list
            The list for corresponding search space.
        """
        candidate_repr = search_space_item["candidates"]
        multihot_list = [False] * mutable.n_candidates
        for i, v in zip(idx, value):
            assert 0 <= i < mutable.n_candidates and candidate_repr[i] == v, \
                "Index '{}' in search space '{}' is not '{}'".format(i, candidate_repr, v)
            assert not multihot_list[i], "'{}' is selected twice in '{}', which is not allowed.".format(i, idx)
            multihot_list[i] = True
        return torch.tensor(multihot_list, dtype=torch.bool)  # pylint: disable=not-callable

    def sample_search(self):
        """
        See :meth:`sample_final`.
        """
        return self.sample_final()

    def sample_final(self):
        """
        Convert the chosen arch and apply it on model.
        """
        assert set(self._chosen_arch.keys()) == set(self._search_space.keys()), \
            "Unmatched keys, expected keys '{}' from search space, found '{}'.".format(self._search_space.keys(),
                                                                                       self._chosen_arch.keys())
        result = dict()
        for mutable in self.mutables:
            if isinstance(mutable, (LayerChoice, InputChoice)):
                assert mutable.key in self._chosen_arch, \
                    "Expected '{}' in chosen arch, but not found.".format(mutable.key)
                data = self._chosen_arch[mutable.key]
                assert isinstance(data, dict) and "_value" in data and "_idx" in data, \
                    "'{}' is not a valid choice.".format(data)
            if isinstance(mutable, LayerChoice):
                result[mutable.key] = self._sample_layer_choice(mutable, data["_idx"], data["_value"],
                                                                self._search_space[mutable.key]["_value"])
            elif isinstance(mutable, InputChoice):
                result[mutable.key] = self._sample_input_choice(mutable, data["_idx"], data["_value"],
                                                                self._search_space[mutable.key]["_value"])
            elif isinstance(mutable, MutableScope):
                logger.info("Mutable scope '%s' is skipped during parsing choices.", mutable.key)
            else:
                raise TypeError("Unsupported mutable type: '%s'." % type(mutable))
        return result

    def _standalone_generate_chosen(self):
        """
        Generate the chosen architecture for standalone mode,
        i.e., choose the first one(s) for LayerChoice and InputChoice.
        ::
            { key_name: {"_value": "conv1",
                         "_idx": 0} }
            { key_name: {"_value": ["in1"],
                         "_idx": [0]} }
        Returns
        -------
        dict
            the chosen architecture
        """
        chosen_arch = {}
        for key, val in self._search_space.items():
            if val["_type"] == LAYER_CHOICE:
                choices = val["_value"]
                chosen_arch[key] = {"_value": choices[0], "_idx": 0}
            elif val["_type"] == INPUT_CHOICE:
                choices = val["_value"]["candidates"]
                n_chosen = val["_value"]["n_chosen"]
                if n_chosen is None:
                    n_chosen = len(choices)
                chosen_arch[key] = {"_value": choices[:n_chosen], "_idx": list(range(n_chosen))}
            else:
                raise ValueError("Unknown key '%s' and value '%s'." % (key, val))
        return chosen_arch

    def random_generate_chosen(self):
        """
        Generate the chosen architecture for standalone mode,
        i.e., choose the first one(s) for LayerChoice and InputChoice.
        ::
            { key_name: {"_value": "conv1",
                         "_idx": 0} }
            { key_name: {"_value": ["in1"],
                         "_idx": [0]} }
        Returns
        -------
        dict
            the chosen architecture
        """
        chosen_arch = {}
        np.random.seed(self.trial_id)
        random.seed(self.trial_id)
        for key, val in self._search_space.items():
            if val["_type"] == LAYER_CHOICE:
                choices = val["_value"]
                
                chosen_idx = np.random.randint(len(choices))
                chosen_arch[key] = {"_value": choices[chosen_idx], "_idx": chosen_idx}
            elif val["_type"] == INPUT_CHOICE:
                choices = val["_value"]["candidates"]
                n_chosen = val["_value"]["n_chosen"]
                if n_chosen is None:
                    n_chosen = len(choices)
                chosen_idx = random.sample(list(range(n_chosen)),n_chosen)
                chosen_arch[key] = {"_value": [choices[idx] for idx in chosen_idx], "_idx": chosen_idx}
            else:
                raise ValueError("Unknown key '%s' and value '%s'." % (key, val))
        return chosen_arch
    

    def _generate_search_space(self):
        """
        Generate search space from mutables.
        Here is the search space format:
        ::
            { key_name: {"_type": "layer_choice",
                         "_value": ["conv1", "conv2"]} }
            { key_name: {"_type": "input_choice",
                         "_value": {"candidates": ["in1", "in2"],
                                    "n_chosen": 1}} }
        Returns
        -------
        dict
            the generated search space
        """
        search_space = {}
        for mutable in self.mutables:
            # for now we only generate flattened search space
            if isinstance(mutable, LayerChoice):
                key = mutable.key
                val = mutable.names
                search_space[key] = {"_type": LAYER_CHOICE, "_value": val}
            elif isinstance(mutable, InputChoice):
                key = mutable.key
                search_space[key] = {"_type": INPUT_CHOICE,
                                     "_value": {"candidates": mutable.choose_from,
                                                "n_chosen": mutable.n_chosen}}
            elif isinstance(mutable, MutableScope):
                logger.info("Mutable scope '%s' is skipped during generating search space.", mutable.key)
            else:
                raise TypeError("Unsupported mutable type: '%s'." % type(mutable))

        return search_space

    def _dump_search_space(self, file_path):
        with open(file_path, "w") as ss_file:
            json.dump(self._search_space, ss_file, sort_keys=True, indent=2)

    def _generate_selected_space(self,file_path):
        with open(file_path, "w") as ss_file:
            json.dump(self._chosen_arch, ss_file, sort_keys=True, indent=2)