【神經(jīng)網(wǎng)絡(luò)搜索】Microsoft NNI 有關(guān)NAS的核心類

【GiantPandaCV導(dǎo)語】本文介紹NNI PyTorch版實現(xiàn)神經(jīng)網(wǎng)絡(luò)過程搜索過程中的幾個重要的類，比如LayerChoice和InputChoice，對這兩個類有了初步認(rèn)識以后，就可以設(shè)計自己的搜索空間。

1. Mutable類

上圖是NNI的有關(guān)NAS的部分類圖，Mutable類表示的意思是可變的，這也是實現(xiàn)NAS中的核心，操作是可變動的，具體選擇什么操作需要優(yōu)化器，也就是tuner來決定。

• Mutable被設(shè)計成一個普通層，具有所有操作的權(quán)重。

• Mutator中應(yīng)該包含網(wǎng)絡(luò)架構(gòu)的狀態(tài)和權(quán)重，而不是層本身。

• Mutable對象有一個key，用于標(biāo)記mutable對象的身份。用戶可以根據(jù)key來進(jìn)行共享不同mutable對象之間的決定。

• 在Mutator的實現(xiàn)中，Mutator應(yīng)該使用key區(qū)分不同的mutable對象。如果兩個mutable對象的key是相同的，說明并不需要對其進(jìn)行區(qū)分，即這兩個mutable對象是相似的。

• 當(dāng)前key的默認(rèn)作用域是全局的。默認(rèn)情況下，key使用counter從1開始計數(shù)，來自動生成unique id

• Mutable類屬于模型級別的設(shè)置，counter是程序級別的。

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

在初始化的時候，需要接收key，如果沒有特別設(shè)置key，那就通過global_mutable_counting()方法返回全局變量counter數(shù)量。

2. MutableScope

MutableScope代碼實現(xiàn)非常短，如下：

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)

MutableScope是繼承了Mutable對象，也有一個key，他是比操作更高層次的抽象。類似的概念有子圖，子模塊，可以看作一系列操作的集合。

• MutableScope可以更好的幫助Mutator做決策，將其看作略高層次的抽象。
• 如果沒有標(biāo)注為mutable scope, 那么搜索空間將會展開為一個列表。如果一個模塊是在mutable scope中定義，那么將被視為sub-search-space, 子搜索空間，并且這些mutable scope之間也可以相互嵌套。
• Mutator有兩種方法使用mutable scope：
• 一種是初始化的時候，通過樹的形式進(jìn)行初始化搜索空間。
• 另一種是實現(xiàn)' enter_mutable_scope '和' exit_mutable_scope '兩個方法
• Mutable Scope也是一種Mutable對象，只不過其比較特殊，包含的內(nèi)容不是普通的操作opration, 而是Mutable對象。Mutable Scope也會在搜索空間中被枚舉出來，但是不應(yīng)該出現(xiàn)在選項的字典中。

3. LayerChoice

LayerChoice類的核心功能是從候選操作中挑選一個，將該操作施加到輸入得到輸出結(jié)果。在特殊情況下，可以選擇zero 或者選擇多個操作。Layer Choice不允許嵌套。主要有以下幾個參數(shù)：

• op_candidates: 候選操作，可以是nn.Module列表或字典
• reduction: 可以從mean, concat, sum, none幾種選擇。
• return_mask: 決定返回結(jié)果是否包含mask
• key：input_choice的key

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) # 添加模塊進(jìn)來
                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

可以看出LayerChoice就是一個類似于列表的類，其中包含了候選的操作，可以通過add_module的方式將候選操作添加到LayerChoice這個類中。

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

前向傳播的時候，是mutator的on_forward_layer_choice函數(shù)進(jìn)行控制具體的操作，return_mask控制是否同時輸出mask和tensor。

一個調(diào)用的例子：

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

4. InputChoice

InputChoice是用來解決網(wǎng)絡(luò)層與層之間連接的問題，有以下幾個參數(shù)：

• n_candidates: 是一個數(shù)，選擇多少個作為input
• choose_from: 是一個裝滿key的列表，都是過去已經(jīng)生成的mutable對象的key。也可以是InputChoice.NO_KEY代表
• n_chosen: 選擇的輸入的個數(shù)，如果不設(shè)置，那就可以選擇任何數(shù)量的組合。
• reduction: 規(guī)約方式有mean, concat, sum, none。
• return_mask&key同上。

綜合來說，Input Choice就是從choose_from對應(yīng)key中選擇n_chosen個輸入， 其中n_candidates決定了forward函數(shù)中，候選選項中選擇的個數(shù)。

舉個例子：

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])

InputChoice的源碼實現(xiàn)：

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

前向傳播的選擇還是通過調(diào)用mutator的on_forward_input_choice函數(shù)來決定選擇哪條路徑連接。

本文主要介紹了nni中搜索空間指定最核心的幾個類，通過使用這些類就可以做到構(gòu)建自己的搜索空間。最近nni更新了2.1版本retiarii等新的功能特性, 允許用戶以高度的靈活性表達(dá)各種搜索空間，重用許多前沿搜索算法，更加易用，準(zhǔn)備踩坑。