单机玩转神经网络架构搜索(NAS) - Auto-Keras学习笔记

介绍

AutoML这个topic在机器学习领域越来越火，新的研究成果也是层出不穷。在网络架构(NAS)，模型压缩(AMC)，数据增强(AutoAugment)，优化器设计(Neural Optimizer Search)，平台相关优化(AutoTVM)等领域，我们都可以看到相应的研究成果表明机器学习可以达到比人肉调参更优的结果。自动化方法正在逐步替代调参工。相信不久的将来，我们面对一个场景，只要喂数据，其他的由机器一条龙完成，而且还能比人肉优化出来人牛X。它们本质上都是大规模空间中的最优参数搜索问题，因此很多方法基本都是相通的。其中网络架构的搜索(Neural Architecture Search, NAS)一直是业界主要研究对象，各种云服务也开始推出相关的服务。一开始的时候，这项技术需要依赖巨大的计算成本。随着各种改进，它的算力需求大大减少，使得它开始有可能飞入寻常百姓家。之前写过一篇简单讨论了相关的技术《神经网络架构搜索（Neural Architecture Search）杂谈》，这里就不重复了。本文主要学习一下Auto-Keras这个开源项目中的实现，使用它我们可以在本地进行网络架构搜索。官方有一篇相关的paper介绍-Auto-Keras: An Efficient Neural Architecture Search System。

安装过程很简单，详见官方介绍。图方便的话直接用pip安装即可：

pip install autokeras

比较核心的搜索过程如下示意图。首先生成几个初始网络架构作为种子并训练得到准确率指标，然后不断通过network morphism产生新的网络架构，这样就形成一个树形的搜索结构。搜索过程中会试图找最有『潜力』(用acquisition function衡量)的网络架构，这个架构会通过训练得到其准确率。这些挑选出来的架构及其准确率会用来拟合高斯过程模型，这个模型又会帮助下一轮树型结构中搜索时预测生成网络架构的准确率指标。

流程

下面沿着官方自带的例子mnist.py来看一下大体实现过程。

    clf = ImageClassifier(verbose=True)
    clf.fit(x_train, y_train, time_limit=12 * 60 * 60)
    clf.final_fit(x_train, y_train, x_test, y_test, retrain=True)
    y = clf.evaluate(x_test, y_test)

因为是图片分类任务，这里先创建ImageClassifier类，接着fit()函数进行网络架构搜索，final_fit()函数对搜索出来的最优模型训练，最后evaluate()函数进行模型评估。其中涉及到的几个主要类的关系：

ImageClassifier继承自ImageSupervised类(实现位于image_supervised.py文件)。它用来在图片分类任务中针对数据集搜索最优的CNN架构。ImageSupervised本身的构造函数很简单，它会调用其父类DeepTaskSupervised(实现位于supervised.py)的构造函数。父类构造函数中会检测是否要继续上次的任务。如果是，则从文件中把分类器和CNN模块恢复回来，否则就新创建CnnModule对象（它的实现在net_module.py文件）。CnnModule构造时会往generators数组里放入三个生成器CnnGenerator, ResNetGenerator和DenseNetGenerator，分别用于生成普通CNN，ResNet和DenseNet。

接下来大头是fit()函数。该函数先计算数据集中size的中值，再把数据都统一resize成这个尺寸。然后就是调用父类DeepTaskSupervised的fit()函数(位于supervised.py文件)。这里边先是准备数据集(切分验证集，数据转换等)，最后调用CnnModule的fit()函数。它本质上是调用其基类NetworkModule的fit()函数。这个函数首次会根据参数创建相应的Searcher对象(默认是BayesianSearcher)。然后进入循环，每次迭代中调用BayesianSearcher的search()函数。在BayesianSearcher的search()函数中：

1. 首次会调用init_search()函数初始化。这里会用到之前的三个网络生成器，分别对于数据集的类别数量和输入shape实例化，并通过它们的generate()函数生成相应的网络。网络计算图用Graph类表示。生成的网络加入到training_queue中等待训练和评估。这个队列中的元素为三元组(graph, other_info, model_id)。对于BayesianSearcher来说other_info就是其父节点id。
2. 从training_queue中弹出一个元素。这个元素代表一个网络结构。调用mp_search()函数(非Colab环境)进行训练和网络的搜索。

mp_search()函数主要有这么几步：

1. 这个函数首先会新建一个进程来作模型的训练(通过train()函数)和评估。这主要是由ModelTrainer的train_model()函数(实现在model_trainer.py文件)完成。由于这里只是为了评估网络架构，没必要训练得很充分，所以会采用early stop节约资源，默认如果5次epoch loss不再减少则提前退出训练。
2. 同时，在当前进程调用_search_common()函数进行网络结构的搜索。这个函数中当发现当前已没有待训练模型时，会首先调用generate()函数生成下一个网络结构。这里就用到BayesianOptimizer了，它会涉及到三个关键数据结构：1) IncrementalGaussianProcess是高斯过程模型用来对网络的指标进行建模和预测；2) SearchTree为搜索树，用来组织搜索过的历史网络结构，A通过network morphism生成B，则A为B的父结点; 3) PriorityQueue将网络结构按其指标进行排序，它用于指导优先从哪些网络结构开始拓展。回到generate()函数。它首先将历史中的网络模型及指标拿出来进行排序，并放入PriorityQueue。。接着进入循环。每次迭代中，从优先级队列中取元素(即指标最好)，并通过模拟退火算法确定是否要对该元素对应的网络进行变形。模拟退火是为了考虑exploration。模拟退火中的acceptance function中的energe function这里也是acquisition function，它某种程度上指示搜索过程中对结点的偏好。因为我们对网络指标的估计不是准确的，还得考虑估计的不确定性，因此这里采用的是UCB。它的计算不仅需要指标预测的均值，还需要方差，这些信息是由高斯过程模型给出的。如果某结点被选定为“潜力股”，那会调用transform()函数(实现在net_transformer.py文件)通过network morphism对其进行扩展。通过变形形成的新网络如果不是重复的，则也会被加入PriorityQueue进行下一轮拓展。这个循环结束后，变量opt_acq、target_graph、和father_id分别记录了最优的acquisition function值，对应的图结构和其父结点ID。这个生成的图在_search_common()中会加入到training_queue中待训练。注意这个生成的图中的权重是可以和父结点重用的(通过load_modoel_by_id加载)，这也是network morphism的最大好处。
3. 从训练的进程拿评估后的指标。如果指标有效，先调用add_model()函数记录相关信息，如将网络图dump到文件，记录到历史，记录最好模型ID等。然后调用update()函数。该函数调用BayesianOptimizer的fit()和add_child()函数。前者主要用于更新高斯过程模型参数；后者主要将该网络结构加入到搜索树当中，以便以后评估是否要对其进行扩展。

以上就是网络搜索的主要过程，大致流程图如下。

例子中，后面分别调用了final_fit()和evaluate()函数进行最优模型的训练和评估。final_fit()函数对前面搜索到的最优网络架构进行权重的训练。在搜索过程中的训练只是为了在候选网络间对比，因此训练得并不充分。这里会将max_no_improvement_num设为30，即30次epoch loss没有减少才退出。最后evaluate()函数将前一步训练好的模型拿来在测试集作预测并得出最终准确率。

实验

MNIST因为数据集太简单，初始时就很高了，貌似不是很能体现出网络架构搜索的价值。我们拿稍微复杂一些的CIFAR10来做下实验。经过在我的低端GPU上搜索约20个小时，可以看到，验证集上准确率从一开始的65%左右到90%以上。最优模型出现在第38轮迭代，验证集准确率为93.16%。基于该模型经过充分训练后测试集上准确率为95.56%。

这里搜索到的最优模型的网络架构如下图。自动搜索出来的网络架构看起来是比较诡异一些。。。