文章转载于:GiantPandaCV
作者: pprp
【GiantPandaCV导语】本文介绍NNI PyTorch版实现神经网络过程搜索过程中的几个重要的类，比如LayerChoice和InputChoice，对这两个类有了初步认识以后，就可以设计自己的搜索空间。
1. Mutable类
上图是NNI的有关NAS的部分类图，Mutable类表示的意思是可变的，这也是实现NAS中的核心，操作是可变动的，具体选择什么操作需要优化器，也就是tuner来决定。
- Mutable被设计成一个普通层，具有所有操作的权重。
- Mutator中应该包含网络架构的状态和权重，而不是层本身。
- Mutable对象有一个key，用于标记mutable对象的身份。用户可以根据key来进行共享不同mutable对象之间的决定。
- 在Mutator的实现中，Mutator应该使用key区分不同的mutable对象。如果两个mutable对象的key是相同的，说明并不需要对其进行区分，即这两个mutable对象是相似的。
- 当前key的默认作用域是全局的。默认情况下，key使用counter从1开始计数，来自动生成unique id
- Mutable类属于模型级别的设置，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，如果没有特别设置key，那就通过global\_mutable\_counting()方法返回全局变量counter数量。
2. MutableScope
MutableScope代码实现非常短，如下：
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做决策，将其看作略高层次的抽象。
- 如果没有标注为mutable scope, 那么搜索空间将会展开为一个列表。如果一个模块是在mutable scope中定义，那么将被视为sub-search-space, 子搜索空间，并且这些mutable scope之间也可以相互嵌套。
- Mutator有两种方法使用mutable scope：
- 一种是初始化的时候，通过树的形式进行初始化搜索空间。
- 另一种是实现' enter\_mutable\_scope '和' exit\_mutable\_scope '两个方法
- Mutable Scope也是一种Mutable对象，只不过其比较特殊，包含的内容不是普通的操作opration, 而是Mutable对象。Mutable Scope也会在搜索空间中被枚举出来，但是不应该出现在选项的字典中。
3. LayerChoice
LayerChoice类的核心功能是从候选操作中挑选一个，将该操作施加到输入得到输出结果。在特殊情况下，可以选择zero 或者选择多个操作。Layer Choice不允许嵌套。主要有以下几个参数：
- op\_candidates: 候选操作，可以是nn.Module列表或字典
- reduction: 可以从mean, concat, sum, none几种选择。
- return\_mask: 决定返回结果是否包含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) # 添加模块进来 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函数进行控制具体的操作，return\_mask控制是否同时输出mask和tensor。
一个调用的例子：
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是用来解决网络层与层之间连接的问题，有以下几个参数：
- n\_candidates: 是一个数，选择多少个作为input
- choose\_from: 是一个装满key的列表，都是过去已经生成的mutable对象的key。也可以是InputChoice.NO\_KEY代表
- n\_chosen: 选择的输入的个数，如果不设置，那就可以选择任何数量的组合。
- reduction: 规约方式有mean, concat, sum, none。
- return\_mask&key同上。
综合来说，Input Choice就是从choose\_from对应key中选择n\_chosen个输入， 其中n\_candidates决定了forward函数中，候选选项中选择的个数。
举个例子：
class Cell(MutableScope): passclass 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的源码实现：
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前向传播的选择还是通过调用mutator的on\_forward\_input\_choice函数来决定选择哪条路径连接。
本文主要介绍了nni中搜索空间指定最核心的几个类，通过使用这些类就可以做到构建自己的搜索空间。最近nni更新了2.1版本retiarii等新的功能特性, 允许用户以高度的灵活性表达各种搜索空间，重用许多前沿搜索算法，更加易用，准备踩坑。
- END -
推荐阅读
更多嵌入式AI技术干货请关注嵌入式AI专栏。
