【神经网络搜索】Efficient Neural Architecture Search¶

class ENASLayer(mutables.MutableScope):

    def __init__(self, key, prev_labels, in_filters, out_filters):
        super().__init__(key)
        self.in_filters = in_filters
        self.out_filters = out_filters

        self.mutable = mutables.LayerChoice([
            ConvBranch(in_filters, out_filters, 3, 1, 1, separable=False),
            ConvBranch(in_filters, out_filters, 3, 1, 1, separable=True),
            ConvBranch(in_filters, out_filters, 5, 1, 2, separable=False),
            ConvBranch(in_filters, out_filters, 5, 1, 2, separable=True),
            PoolBranch('avg', in_filters, out_filters, 3, 1, 1),
            PoolBranch('max', in_filters, out_filters, 3, 1, 1),
            SEConvBranch(in_filters, out_filters, 3, 1, 1, reduction=4)
        ])
        if len(prev_labels) > 0:
            self.skipconnect = mutables.InputChoice(
                choose_from=prev_labels, n_chosen=None)
        else:
            self.skipconnect = None
        self.batch_norm = nn.BatchNorm2d(out_filters, affine=False)

    def forward(self, prev_layers):
        out = self.mutable(prev_layers[-1])
        if self.skipconnect is not None:
            connection = self.skipconnect(prev_layers[:-1])
            if connection is not None:
                out += connection
        return self.batch_norm(out)

class GeneralNetwork(nn.Module):
    def __init__(self, num_layers=6, out_filters=12, in_channels=3, num_classes=10,
                 dropout_rate=0.0):
        super().__init__()
        self.num_layers = num_layers
        self.num_classes = num_classes
        self.out_filters = out_filters
        self.dropout_rate = dropout_rate

        self.stem = nn.Sequential(
            nn.Conv2d(in_channels, out_filters, 3, 1, 1, bias=False),
            nn.BatchNorm2d(out_filters)
        )

        pool_distance = self.num_layers // 3
        # 进行pool操作是num_layers // 3
        self.pool_layers_idx = [pool_distance - 1, 2 * pool_distance - 1]
        self.dropout = nn.Dropout(self.dropout_rate)

        self.layers = nn.ModuleList()  # convolutional
        self.pool_layers = nn.ModuleList()  # reduction

        labels = []
        for layer_id in range(self.num_layers):  # 设置12个layer
            labels.append("layer_{}".format(layer_id))

            if layer_id in self.pool_layers_idx:  # 如果使用pool
                self.pool_layers.append(FactorizedReduce(
                    self.out_filters, self.out_filters))

            self.layers.append(  # 相当于Node节点
                ENASLayer(labels[-1], labels[:-1], self.out_filters, self.out_filters))

        self.gap = nn.AdaptiveAvgPool2d(1)
        self.dense = nn.Linear(self.out_filters, self.num_classes)

    def forward(self, x):
        bs = x.size(0)
        cur = self.stem(x)  

        layers = [cur]

        for layer_id in range(self.num_layers):
            cur = self.layers[layer_id](layers)
            layers.append(cur)
            if layer_id in self.pool_layers_idx:
                # 如果轮到了池化层
                for i, layer in enumerate(layers):
                    layers[i] = self.pool_layers[self.pool_layers_idx.index(
                        layer_id)](layer)
                cur = layers[-1]

        cur = self.gap(cur).view(bs, -1)
        cur = self.dropout(cur)
        logits = self.dense(cur)
        return logits

model = GeneralNetwork()
trainer = enas.EnasTrainer(model,
                           loss=criterion,
                           metrics=accuracy,
                           reward_function=reward_accuracy,
                           optimizer=optimizer,
                           callbacks=[LRSchedulerCallback(lr_scheduler), ArchitectureCheckpoint("./checkpoints")],
                           batch_size=args.batch_size,
                           num_epochs=num_epochs,
                           dataset_train=dataset_train,
                           dataset_valid=dataset_valid,
                           log_frequency=args.log_frequency,
                           mutator=mutator)

def _sample_layer_choice(self, mutable):
    # 选择 某个层 只需要选一个就可以了
    self._lstm_next_step() # 让_inputs在lstm中进行一次前向传播

    logit = self.soft(self._h[-1]) # linear 从隐藏层embedd得到可选的层的逻辑评分

    if self.temperature is not None:
        logit /= self.temperature # 一个常量 貌似是RL中的trick

    if self.tanh_constant is not None:
        # tanh_constant * tanh(logits) 用tanh再激活一次（可选）
        logit = self.tanh_constant * torch.tanh(logit)

    if mutable.key in self.bias_dict:
        logit += self.bias_dict[mutable.key]
        # 对卷积层进行了偏好处理，如果是卷积层，那就在对应的值加上一个0.25，增大被选中的概率

    # softmax, view(-1), 
    branch_id = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1) 
    # 依据概率来选下角标，如果数量不为1，选择的多个中没有重复的 
    # eg: [100,1,1] 最有可能选择100对应的下标0

    log_prob = self.cross_entropy_loss(logit, branch_id) # 交叉熵损失函数 - 判断logit和branchid分布是否相似程度

    self.sample_log_prob += self.entropy_reduction(log_prob) # 求和或者求平均

    entropy = (log_prob * torch.exp(-log_prob)).detach()  # pylint: disable=invalid-unary-operand-type　??

    self.sample_entropy += self.entropy_reduction(entropy) # 样本熵？

    self._inputs = self.embedding(branch_id) # 得到对应id的embedding, 从选择空间 - 映射到 - 隐空间

    return F.one_hot(branch_id, num_classes=self.max_layer_choice).bool().view(-1) # 将选择变成one_hot向量

if self.entropy_weight: # 交叉熵权重 
    reward += self.entropy_weight * self.mutator.sample_entropy.item() # 得到样本熵