【神经网络搜索】Microsoft NNI 有关NAS的核心类

class Mutable(nn.Module):
    def __init__(self, key=None):
        super().__init__()
        if key is not None:
            if not isinstance(key, str):
                key = str(key)
                logger.warning("Warning: key \"%s\" is not string, converted to string.", key)
            self._key = key
        else:
            self._key = self.__class__.__name__ + str(global_mutable_counting())
        self.init_hook = self.forward_hook = None

class MutableScope(Mutable):
    def __init__(self, key):
        super().__init__(key=key)

    def __call__(self, *args, **kwargs):
        try:
            self._check_built()
            self.mutator.enter_mutable_scope(self)
            return super().__call__(*args, **kwargs)
        finally:
            self.mutator.exit_mutable_scope(self)

class LayerChoice(Mutable):
    def __init__(self, op_candidates, reduction="sum", return_mask=False, key=None):
        super().__init__(key=key)
        self.names = []
        if isinstance(op_candidates, OrderedDict):
            for name, module in op_candidates.items():
                assert name not in ["length", "reduction", "return_mask", "_key", "key", "names"], \
                    "Please don't use a reserved name '{}' for your module.".format(name)
                self.add_module(name, module) # 添加模块进来
                self.names.append(name)

        elif isinstance(op_candidates, list): 
            for i, module in enumerate(op_candidates):
                self.add_module(str(i), module)
                self.names.append(str(i)) # list的画就手动添加name
        else:
            raise TypeError("Unsupported op_candidates type: {}".format(type(op_candidates)))
        self.reduction = reduction
        self.return_mask = return_mask # 是否同时return mask 和 tensor

def forward(self, *args, **kwargs):
    """
    Returns
    -------
    tuple of tensors
        Output and selection mask. If ``return_mask`` is ``False``, only output is returned.
    """
    out, mask = self.mutator.on_forward_layer_choice(self, *args, **kwargs)
    if self.return_mask:
        return out, mask
    return out

self.op_choice = LayerChoice(OrderedDict([
    ("conv3x3", nn.Conv2d(3, 16, 128)),
    ("conv5x5", nn.Conv2d(5, 16, 128)),
    ("conv7x7", nn.Conv2d(7, 16, 128))
]))

class Cell(MutableScope):
    pass

class Net(nn.Module):
    def __init__(self):
        self.cell1 = Cell("cell1")
        self.cell2 = Cell("cell2")
        self.op = LayerChoice([conv3x3(), conv5x5()], key="op")
        self.input_choice = InputChoice(choose_from=["cell1", "cell2", "op", InputChoice.NO_KEY])

    def forward(self, x):
        x1 = max_pooling(self.cell1(x))
        x2 = self.cell2(x)
        x3 = self.op(x)
        x4 = torch.zeros_like(x)
        return self.input_choice([x1, x2, x3, x4])

class InputChoice(Mutable):
    NO_KEY = ""

    def __init__(self, n_candidates=None, choose_from=None, n_chosen=None,
                 reduction="sum", return_mask=False, key=None):
        super().__init__(key=key)
        # precondition check
        assert n_candidates is not None or choose_from is not None, "At least one of `n_candidates` and `choose_from`" \
                                                                    "must be not None."
        if choose_from is not None and n_candidates is None:
            n_candidates = len(choose_from) # choose_from 不为None，n_candidate就是其长度
        elif choose_from is None and n_candidates is not None:
            choose_from = [self.NO_KEY] * n_candidates # 将空白字符串作为key
        assert n_candidates == len(choose_from), "Number of candidates must be equal to the length of `choose_from`."
        assert n_candidates > 0, "Number of candidates must be greater than 0."
        assert n_chosen is None or 0 <= n_chosen <= n_candidates, "Expected selected number must be None or no more " \
                                                                  "than number of candidates."

        self.n_candidates = n_candidates
        self.choose_from = choose_from.copy()
        self.n_chosen = n_chosen
        self.reduction = reduction
        self.return_mask = return_mask

    def forward(self, optional_inputs):
        # optional_inputs是一个列表，里边是所有可选的输入张量
        optional_input_list = optional_inputs
        
        if isinstance(optional_inputs, dict):
            optional_input_list = [optional_inputs[tag] for tag in self.choose_from]
            
        assert isinstance(optional_input_list, list), \
            "Optional input list must be a list, not a {}.".format(type(optional_input_list))
            
        assert len(optional_inputs) == self.n_candidates, \
            "Length of the input list must be equal to number of candidates."
            
        out, mask = self.mutator.on_forward_input_choice(self, optional_input_list)
        if self.return_mask:
            return out, mask
        return out