[论文评析-工程]MAE-DET: Revisiting Maximum Entropy Principle in Zero-Shot NAS for Efficient Object Detecti

文章信息

题目:MAE-DET: Revisiting Maximum Entropy Principle in Zero-Shot NAS for Efficient Object Detection,
发表:ICML 2022,
作者: Zhenhong Sun * 1 Ming Lin * 1 Xiuyu Sun 1 Zhiyu Tan 1 Hao Li 1 Rong Jin 1,

注: 该工作被应用于DAMO-yolo中用于搜索检测网络的backbone, 效果超过了经典的backbone.

背景

当前已有的用于目标检测的NAS方法由于消耗GPU资源并且耗时从而阻碍了其进一步发展和使用. 基于此,作者提出了一种新颖的zero-shot NAS方法, 该方法基于最大熵原则自动高效的搜索最优的backbone而无需训练神经网络的参数. 结果表明,在跟定约束条件下,搜索出来的backbone网络性能超过了手工设计的backbone.

文章理论性比较强, 这里不深入介绍,主要是结合代码来熟悉以下NAS方法在实际中是怎么回事,以及怎么用的.

项目介绍

项目的组织结构示意图如下:
其中:
configs/下为配置文件,包含各个模块的参数配置
latency/下为作者实现的latency计算模块
script/下为特定任务的执行脚本
nas/下为系统的核心, 其包含的各个模块的逻辑关系如下图所示:

builder.py模块主要是定义masternet和Score, Latency, Search space三大模块,这四块也就是所谓扼的NAS model.

search.py主要时初始化NAS model,并初始化Population module,然后选择-变异-进化来搜索最优的backbone,

那么系统如下运行呢? 下面以通过NAS搜索最佳的small backbone为例, scripts/damo-yolo/example_k1kx_small.sh打开后可以看到:
先准备环境创建工作目录work_dir, 然后把自定义的初始化backbone写入到work_dir/init_structure.txt文件中, 这个文件定义了ZNS搜索backbone的起点, 通过space_structure_txt参数传入, 在后面会用到.
接下来就是通过mpirun来执行nan/search.py来搜索backbone了.

注: 相关参数配置见configs/configure_nas.py的默认配置, 另外也可以通过–cfg_options来直接指定一些参数的值.

下面来看下search.py文件.
可以清楚的看到系统的初始化过程.

接下来利用之前提到过的space_structure_txt参数指定的init_structure.txt文件来真正的初始化MasterNet,并获取网络相关信息来初始化Population,
注: 由于网络的backbone一直在不断进化, 因此不能写死, MasterNet定义的初始化方式是通过属性structure_info来初始化网络, 后面每次backbone更新后只需要获取更新后的structure_info并重新初始化网络即可.

上述步骤完成后,接下来就是NAS的核心部分: 执行do_main_job, do_main_job在干啥?实际上就是迭代执行-选择-变异-进化,然后评估来搜索最佳的backbone,

下面来详细看下do_main_job.

可以看到每次通过random.choice从population中随机选择一个backbone结构作为初始的bakbone (init_random_structure_info),

然后通过get_new_random_structure_info方法变异产生新的backbone (random_structure_info), 然后利用新产生的random_structure_info执行get_info_for_evaluation来再次初始化MasterNet, 获取相关信息random_structure_info后执行update_population方法更新Population.

下面来详细看下get_new_random_structure_info中是如何进行变异的,
可以看到该方法内部实际上就是在调用mutation_function方法进行变异.

那么这个mutate_function到底是啥?
其取值为model_nas.mutation, 它的具体含义是什么?
可以看到首先依据cfg.space_mutation (可以使用config_nas.py中的默认值,也可以在执行脚本sh中指定, 这里的值是space_K1KXK1, 为搜索空间标识,对应nas/spaces下的相应python文件space_K1KXK1.py)确定编译函数的路径,然后通过load_py_module_from_path方法加载对应python文件下的mutation函数对象作为mutation属性的值. 简单一句话来说就是加载对应搜索空间中的mutate_function函数.

下面来看下space_K1KXK1.py中的mutate_function的实现:
mutate_function是针对特定架构实现的, 首先该文件中定义了k1kxk1架构下每个block变异的方法(channel, kernel_size, layer_number等)以及每个方法对应参数可调的范围,
可以看到K1KXK1对应的两种类型的block: ‘ConvKXBNRELU’, ‘SuperResConvK1KX’, 对于指定的block, 根据其class_name进入到相应的流程中后, 首先通过random.choice随机选择变异方法, 然后执行相应的方法进行变异,
值得注意的是:变异时可以根据需要人工施加一些约束, 这些约束定义在.space_K1KXK1.py的开头.

def mutate_function(block_id, structure_info_list, budget_layers, minor_mutation=False):

    structure_info = structure_info_list[block_id]
    if block_id < len(structure_info_list)-1: 
        structure_info_next = structure_info_list[block_id+1]
    structure_info = copy.deepcopy(structure_info)
    class_name = structure_info['class']

    if class_name == 'ConvKXBNRELU':
        if block_id <= len(structure_info_list) - 2:
            random_mutate_method = random.choice(stem_mutate_method_list)
        else:
            return False

        if random_mutate_method == 'out':
            new_out = mutate_channel(structure_info['out'])
            # Add the constraint: the maximum output of the stem block is 128
            new_out = min(32, new_out)
            if block_id < len(structure_info_list)-1:
                new_out = min(structure_info_next['out'], new_out)
            structure_info['out'] = new_out
            return [structure_info]

        if random_mutate_method == 'k':
            new_k = mutate_kernel_size(structure_info['k'])
            structure_info['k'] = new_k
            return [structure_info]

    elif class_name == 'SuperResConvK1KX':
        # coarse2fine mutation flag, only mutate the channels' output
        mutate_method_list_final=['out', 'btn'] if minor_mutation else mutate_method_list

        random_mutate_method = random.choice(mutate_method_list_final)

        if random_mutate_method == 'out':
            new_out = mutate_channel(structure_info['out'])
            # Add the contraint: output_channel should be in range [min, max]
            if channel_range[block_id] is not None:
                this_min, this_max = channel_range[block_id]
                new_out = max(this_min, min(this_max, new_out))
            # Add the constraint: output_channel > input_channel
            new_out = max(structure_info['in'], new_out)
            if block_id < len(structure_info_list) - 1:
                new_out = min(structure_info_next['out'], new_out)
            structure_info['out'] = new_out
            if block_id < len(structure_info_list) - 1 and "btn" in structure_info_next:
                new_btn = min(new_out, structure_info_next['btn'])
                structure_info_next['btn'] = new_btn

        if random_mutate_method == 'k':
            new_k = mutate_kernel_size(structure_info['k'])
            structure_info['k'] = new_k

        if random_mutate_method == 'btn':
            new_btn = mutate_channel(structure_info['btn'])
            # Add the constraint: bottleneck_channel <= output_channel
            new_btn = min(structure_info['out'], new_btn) 
            structure_info['btn'] = new_btn

        if random_mutate_method == 'L':
            new_L = mutate_layer(structure_info['L'])
            # add the constraint: the block 1 can't have the large layers.
            if block_id==1:
                new_L = min(2, new_L)
            else:
                new_L = min(int(budget_layers//2//(len(structure_info_list)-2)), new_L)

            structure_info['L'] = new_L

        # add the constraint: the btn must be larger than out/btn_minimum_ratio.
        if structure_info['btn']<(structure_info['out']/btn_minimum_ratio):
            structure_info['btn'] = smart_round(structure_info['out']/btn_minimum_ratio)
        
        return [structure_info]
    
    else:
        raise RuntimeError('Not implemented class_name=' + class_name)

好了,至此应该对NAS系统整个过程有了直观的感受, 下面来看下论文中的几个核心部分的计算.

核心部分的计算以及实现

论文最核心的在于单次前向传播中用熵来衡量模型的性能.
该行代码指定了score的计算方式: ‘entropy’: ComputeEntropyScore,
self.compute_score = __all_scores__[self.cfg.score_type](self.cfg, logger=self.logger)

每次调用如下方法时会计算score的值.

 random_struct_info = model_nas.get_info_for_evolution(structure_info=random_structure_info)

其中model_info["score"] = self.do_compute_nas_score(model)

再来看do_compute_nas_score方法, 最终会定位到nas/scores/ComputeEntropyScore.py中,
可以看到, 获取前向传播输出前的值的方差,然后取对数. 对于不同的架构, 前向传播计算熵略有不同, 其中前向传播的实现位于相应文件(比如SuperResConvK1KXK1.py)下的entropy_forward方法.

总结

自动搜索最佳的Backbone, 只需前向传播, 而无需训练求模型参数, 感觉很神奇.

需要一些优化,并行计算的背景. 工程的要求比较高,好些概念,好些包第一次听到,

代码模块的组织值得学习.

References

1.Sun, Zhenhong, et al. “Mae-det: Revisiting maximum entropy principle in zero-shot nas for efficient object detection.” International Conference on Machine Learning. PMLR, 2022.