云原生自动机器学习系统 Katib 的设计与实现

今天，2019 世界人工智能大会在上海开幕，数百位全球 AI 界领军科学家、产业界领袖齐聚黄浦江畔，探索人工智能、机器学习、深度学习在企业、社会中的应用与实践。

为了成为国内人工智能应用的第一梯队企业，利用 AI 技术消除垂直行业痛点。近年来，越来越多公司开始尝试利用 Kubernetes 来管理机器学习工作负载。相应的，Kubernetes 原生机器学习库 Kubeflow 也受到了更多关注——社区各个项目的贡献者、用户数量持续上升。

在 Kubeflow 社区，Katib 作为云原生自动机器学习系统，近期发布了新版本。为了帮助更多工程师掌握这个云原生自动机器学习系统的底层逻辑，在本文中，才云科技软件工程师 gaocegege 将详细介绍新版本 v1alpha2 的设计与实现。

01 问题背景

如今，机器学习技术和算法几乎应用于每个领域，而建立一个高质量的机器学习模型是一个迭代的、复杂的、耗时的过程：除了要求数据科学家具备高效调整超参数的经验和能力，他们还需要不断尝试不同的算法和技术。

事实上，调参非常考验工程师的知识扎实水平和实践经验。随着数据量的持续爆炸性增长，企业逐渐意识到单纯依靠人力是永远解决不了调参难题的，建立良好的机器学习模型过程自动化至关重要。

目前，深度学习领域有超参数搜索和模型结构搜索等技术，可以用算力来减少人力需求。这些技术的主要目的是通过扮演领域专家的角色，减少人工在该过程中的作用，填补非专家机器学习用户的空白。

而在云原生领域，kubeflow/katib 提供了更友好的超参数搜索与模型结构搜索能力，允许运维工程师管理整个调参过程，降低企业 AI 进入门槛。

02 相关工作

本节介绍了学术界和工业界在超参数搜索领域的一些进展。

超参数搜索

超参数训练在工业界和学术界都受到相当程度的关注，目前也涌现了一批相对成熟的解决方案。

所谓超参数训练，它可以被形式化地表示为在给定多维搜索空间下的黑盒优化问题：对于单一超参数，它可能是离散的值，如激活函数的选择等，也有可能是连续的值；由于超参数数量不确定，它的搜索空间是多维的；而对于超参数训练，其输入参数与输出是相互独立的变量，因此属于黑盒优化的范畴。

当前，工业界较常用的超参数训练算法主要有网格搜索、随机搜索和贝叶斯优化等。

网格搜索是指在所有候选参数中，通过循环遍历尝试每一种的可能性，表现最好的参数就是最终结果。它的思路及实现方式都很简单，但经过笛卡尔积组合后搜索空间会大幅扩大，面对大数据和多参数往往耗时过久，因此只适用于超参数数量小的情况。

随机搜索是 Bengio 等人针对网格搜索的不足提出的另一种方法，它首先为每类超参数定义一个边缘分布，通常取均匀分布，然后在这些参数上采样进行搜索。由于存在随机因素，这种算法的搜索结果可能特别差，也可能特别好。总体来说，它的效率比网格搜索高，但是不保证一定能找到比较好的超参数。

贝叶斯优化是一个经典的黑盒优化算法，它利用了无限维的高斯过程来模拟黑盒的超参数搜索目标函数形式。限于篇幅，此处我们不再赘述，其大致原理是利用之前探索的结果，探索更有可能的新的参数选择。除此之外，学术界还有许多其他的超参数搜索算法，如 HyperBand 等。

模型结构搜索

模型结构搜索是目前学术界的研究热点之一。现阶段，神经网络的结构设计需要特定领域的算法科学家进行人工分析与建模，而模型结构搜索希望通过强化学习或者进化算法的方式，自动寻找到最优的模型结构。

在这一领域，目前比较有影响力的工作有 ENAS、Auto-Keras、DARTS、ProxylessNAS 等。目前的 NAS 都会采用某种意义上的参数共享，来加速模型搜索的过程，进而缓解其对算力的巨大要求。

自动机器学习系统

超参数搜索与模型结构搜索一定程度上是可以通用的算法，因此目前业界也有许多自动机器学习系统，帮助算法科学家们更方便地实现自动机器学习。其中 Google Vizier 对超参数搜索问题进行了非常好的抽象，并且基于此实现了谷歌内部的超参数搜索系统。

Google Vizier 架构图

受这一工作启发，工业界出现了许多开源实现，其中最具代表性的是 advisor。

Auto-Keras 是一个专门用来进行模型结构搜索的库。它提供了模型结构搜索的基准实现，包括随机搜索、网格搜索、贪心算法和贝叶斯优化等。与此同时，它允许用户自行定义新的模型结构搜索算法，并且集成在 Auto-Keras 中。

除此之外，NNI 是由微软亚研院开源的自动机器学习工具。它是既支持模型结构搜索，也支持超参数搜索的工具。它为模型结构搜索实现了一个统一的接口，相比于其他的工具，具有更好的易用性。

Katib 的设计与实现

虽然目前已经有了许多可以帮助算法科学家们快速进行参数搜索的工具，但这些工具大多会引入额外的组件，增加运维负担。在这个背景下，Katib 应运而生。

Katib 旨在实现一个云原生超参数搜索与模型结构搜索系统，在复用 Kubernetes 对 GPU 等资源的管理能力的同时，保证系统的可扩展性。由于是云原生实现，维护这一系统的工作对运维工程师更加友好。

目前的设计

为了实现云原生的架构目标，新版本的 Katib 引入了两个 CRD：Experiment 和 Trial。

Experiment 的定义类似于 Google Vizier 中的 Study，在 Katib v1alpha1 中，它的名字是 StudyJob。它存储了超参数搜索或模型结构搜索的搜索空间，以及 Trial 训练任务的模板和优化的目标值等内容。

Trial 是 Experiment 在确定了搜索空间中的一组参数取值后，进行的模型训练任务。它包含训练任务实例、训练的指标等。当 Trial 的状态为 Completed 时，就意味着这一组参数被验证了，而且对应的训练指标被收集起来了。

而 Job 则是 Kubernetes 上的一个训练实例，它是用来验证一个 Trial 中的参数的。Job 可以是 TensorFlow 模型训练任务 TFJob，可以是 PyTorch 训练任务 PyTorchJob，也可以是以 Kubernetes batchv1/job 运行的任意框架的训练代码。

Suggestion 是在 Experiment 中定义的搜索空间中寻找一组参数取值的算法。寻找到的参数取值会交给 Trial 去运行。目前 Kaitb 支持：

• random search

• grid search

• hyperband

• bayesian optimization

• NAS based on reinforcement learning

• NAS based on EnvelopeNets

Case Study：随机超参数搜索任务

接下来，我们以一次超参数搜索任务为例，从用户的角度看下 katib 的使用流程。

对于用户而言，你只需要创建出 Experiment，即可进行超参数搜索训练：

这一配置包含了完成一个超参数训练所需的所有信息，如需要优化的目标指标、需要用到的搜索算法、Trial 的训练代码，以及搜索空间等。下面我们会一个个介绍。

首先是关于并行相关的配置：

• parallelTrialCount 定义了允许的并行数，在这个例子中，用户允许同一时间内有 3 个运行的 Trial；

• maxTrialCount 是最大的 Trial 数量，在我们的例子中为 12，也就是在 12 个 Trial 运行完后，Experiment 就被标记为完成；

• maxFailedTrialCount 是最大允许失败的次数，在这一例子中，用户可以容忍 3 个 Trial 训练失败。

objective 这一部分是对训练目标的定义。在这一例子中，用户定义了训练的目标是准确率，训练目标是最大化准确率，当准确率到达 0.99 时，超参数训练任务可以被认为已经完成。

algorithm 一项定义了超参数搜索的算法，这里用户希望使用随机搜索。而后面的 trialTemplate 则是定义了 Trial 进行训练时需要的代码。在这一例子中，用户使用了 batch/v1 Job 类型进行训练，使用的框架是 MXNet。

其中我们可以看到模板定义 .HyperParameters，这一定义会在确定了参数取值后被解释执行，生成完整的训练代码定义。

最后，用户定义了需要搜索的参数空间。

在这个例子中，用户一共定义了三个需要搜索的参数，分别是 --lr、--num-layers 和 --optimizer。

其中第一个参数，--lr（也就是 Learning Rate）的搜索空间是从 0.01 到 0.03 的浮点数范围。第二个参数的搜索空间是从 2 到 5 的整形数空间，也就是有 2、3、4、5 四个选择。第三个参数的搜索空间是一个列表，有三个选择，分别是 sgd、adam 和 ftrl。

用户在 Kubernetes 中创建了这一 Experiment 之后，Experiment 会根据用户的定义创建出对应的 Trials：

在每一个 Trial 中，配置如下：

其中：

• metricsCollectorSpec 是一个用于收集训练指标的 Job 的定义，后面会展开介绍；

• Objective 跟 Experiment 中的 Objective 一致；

• parameterAssignments 是 Experiment 定义的搜索空间下的一组参数取值；

• --lr 的取值是 0.016681020412579572，--num-layers 的取值是 3，--optimizer 的取值是 sgd；

• runSpec 根据 parameterAssignments 以及 Experiment 中的 TrialTemplate 定义生成，是真正可以被运行的 Kubernetes 资源定义。

在这一例子中，RunSpec 会被以 Kubernetes batchv1 Job 的方式运行起来，其真实配置如下：

这一 Job 是真正的训练代码，超参数已经通过命令行参数的方式传入到了代码中，而后，训练代码就可以直接利用提供的超参数进行模型训练。

训练任务指标会由之前在 metricsCollectorSpec 中定义的另一个 Job 完成。任务的 Job Name 以命令行参数的方式传入 metricsCollector Job，metricsCollector Job 会利用 Kubernetes Client 获取训练任务的日志，并从日志中处理得到训练的指标，再通过 Katib Manager 存入 Katib DB。

这些训练指标会在后续的超参数搜索算法（如贝叶斯优化）中被用来寻找新的参数取值。

处理流程

在了解 Katib 的处理流程之前，我们先介绍下 Katib 目前有哪些组件：

• Experiment Controller，提供对 Experiment CRD 的生命周期管理；

• Trial Controller，提供对 Trial CRD 的生命周期管理；

• Suggestions，以 Deployment 的方式部署，用 Service 方式暴露服务，提供超参数搜索服务。目前有随机搜索、网格搜索、贝叶斯优化等；

• Katib Manager，一个 GRPC server，提供了对 Katib DB 的操作接口，同时充当 Suggestion 与 Experiment 之间的代理；

• Katib DB，数据库。其中会储存 Trial 和 Experiment，以及 Trial 的训练指标。目前默认的数据库为 MySQL。

Katib 架构图

当一个 Experiment 被创建的时候，Experiment Controller 会先通过 Katib Manager 在 Katib DB 中创建一个 Experiment 对象，并且打上 Finalizer 表明这一对象使用了外部资源（数据库）。

随后，Experiment Controller 会根据自身的状态和关于并行的定义，通过 Katib Manager 提供的 GRPC 接口，让 Manager 通过 Suggestion 提供的 GRPC 接口获取超参数取值，然后再转发给 Experiment Controller。

在这个过程中，Katib Manager 是一个代理的角色，它代理了 Experiment Controller 对 Suggestion 的请求。拿到超参数取值后，Experiment Controller 会根据 TrialTemplate 和超参数的取值，构造出 Trial 的定义，然后在集群中创建它。

Trial 被创建后，与 Experiment Controller 类似，Trial Controller 同样会通过 Katib Manager 在 Katib DB 中创建一个 Trial 对象。随后构造出期望的 Job（如 batchv1 Job、TFJob、PyTorchJob 等）和 Metrics Collector Job，并在集群上创建出来。

这些 Job 运行结束后，Trial Controller 会更新 Trial 的状态，Experiment Controller 会更新 Experiment 的状态。然后 Experiment 会继续下一轮的迭代。之前的 Trial 已经被训练完成，而且训练的指标已经被收集起来了。Experiment 会根据配置，判断是否要再创建新的 Trial，如果需要则再重复之前的流程。

未来规划

目前 Katib v1alpha2 的架构还不是特别完善，它存在一些问题：

• Katib Manager 作为中间代理，有些多余；

• Katib DB 中存在冗余的信息；

• Suggestion 没有被较好地管理起来，目前依赖用户手动维护；

• 模型结构搜索较弱；

• 训练指标收集的方式扩展性较差，而且存在权限问题。

为了解决这些问题，社区目前在做的工作包括：

• 利用 Sidecar 的模式，替代现行训练指标收集的方案；

• 引入新的 CRD：Suggestion；

• 引入统一的数据存储接口，支持不同的数据库。

Katib 架构图 (2)

上图是社区规划中的架构，目前这些工作还在进行中，所以可能会有变动。

新架构中存在三个 CRD：Experiment、Trial 和 Suggestion。同时三个 Controller 不会依赖外部的数据库，不再会在数据库中维护 Experiment 和 Trial 的备份。只有训练指标会在 Metrics Collector Sidecar 中被收集后存入数据库中。

这些指标的最终值会被某些超参数搜索算法，如贝叶斯优化等，用于产生新的超参数取值。这些指标在训练过程中的中间值，则会被 EarlyStopping 服务用于停止训练效果不佳的 Trial。

总结与分析

最后，我们再总结一下 Katib 与 advisor 和 NNI 的相似之处以及不同之处。

如下表所示，Google Vizier 也在列，但由于 Vizier 并不是开源产品，所以我们不对它做具体展开。

同：Katib、advisor 和 NNI 都支持并行地进行超参数搜索。

异：Katib 支持利用 Kubernetes 的集群资源进行超参数搜索；advisor 不支持利用集群资源；NNI 支持不同的平台，如 Microsoft PAI、原生 Kubernetes、支持 Kubeflow 的 Kubernetes 等。

同：支持贝叶斯优化、HyperBand、随机搜索和网格搜索等超参数算法。

异：Katib 目前仅支持以上四种算法，而 advisor 和 NNI 都有更多算法支持。这也是 Katib 需要后续完善的一方面。Katib、advisor 和 NNI 对算法的扩展性都不错。对于早期停止策略的支持，NNI 具有更完善的支持；对于模型结构搜索，NNI 支持最好，Katib 次之，advisor 并无对应功能。

同：具备训练指标收集功能，用于指导后续搜索。

异：NNI 实现了一个 SDK，允许用户通过 SDK 汇报训练指标，属于 Push-based 的实现，具有一定的侵入性。而 Katib 和 advisor 都希望通过 Pull-based 的实现，不修改用户代码，而是分析输出日志。但 Katib 希望未来能同时实现这两种不同的方式，既提供更好的扩展性，又将选择权交给用户。

同：具备超参数注入训练代码功能，支持探索性训练。

异：NNI 无论是超参数搜索还是模型结构搜索，都需要用户在代码中添加 comment 或者利用 SDK 的方式进行改动，具有一定的侵入性。但它也具有一些 AST 级别的 Trick，使得这部分修改在没有 NNI 的情况下，也不会影响正常的执行。

而 Katib 和 advisor 都希望做到对用户代码的零侵入性。Katib 目前在训练指标收集上可以通过日志收集的方式做到零侵入性。而在超参数注入的问题上，后续会探索利用与 NNI 类似的 AST 重写的方式，在运行时改写用户代码，而不再需要目前命令行参数的修改方式。

同：云原生支持。

异：Katib 是 Kubernetes Native 的系统，它的目标就是能够以 Kubernetes Native 的方式部署在 Kubernetes 上并且方便运维。所以 Katib 能够尽可能以 CRD 等 Kubernetes 扩展功能的方式实现自身逻辑，这样的理念使得它对 Kubernetes 和 Kubeflow 的用户非常友好。advisor 是一个独立的系统。NNI 也是一个相对独立的系统，只是它的训练任务可以运行在不同平台上。

同：开源项目。

异：Katib 和 advisor 都是采取 Apache 2.0 许可，NNI 是在 MIT 许可下开源的。

总而言之，当下，如果不考虑云原生部署方式，或者目前对模型结构搜索具有更强的需求，同时不在意对代码添加注释这样的修改方式，这种情况下企业可以选择 NNI。

而如果企业希望自己的超参数搜索与模型结构搜索系统是云原生的，同时非常注重代码的零侵入性，Katib 是更好的选择。

作为 Kubeflow 社区的重要贡献者，才云在这里感谢包括 gaocegege 在内的所有贡献者的工作，同时，我们也欢迎更多贡献者加入社区，一起为完善 Katib 做出更多努力。如果读者对 Katib 感兴趣，或者有疑问，欢迎留言讨论！

参考文献

[1]Elshawi, Radwa, Mohamed Maher, and Sherif Sakr. “Automated Machine Learning: State-of-The-Art and Open Challenges.” arXiv preprint arXiv:1906.02287 (2019).

[2]Bergstra, James, and Yoshua Bengio. “Random search for hyper-parameter optimization.” Journal of Machine Learning Research 13.Feb (2012): 281-305.‏

[3]Pelikan, Martin, David E. Goldberg, and Erick Cantú-Paz. “BOA: The Bayesian optimization algorithm.” Proceedings of the 1st Annual Conference on Genetic and Evolutionary Computation-Volume 1. Morgan Kaufmann Publishers Inc., 1999.

[4]Hyperparameter tuning for machine learning models.

[5]Pham, Hieu, et al. “Efficient neural architecture search via parameter sharing.” arXiv preprint arXiv:1802.03268 (2018).

[6]Jin, Haifeng, Qingquan Song, and Xia Hu. “Efficient neural architecture search with network morphism.” arXiv preprint arXiv:1806.10282 (2018).

[7]Liu, Hanxiao, Karen Simonyan, and Yiming Yang. “Darts: Differentiable architecture search.” arXiv preprint arXiv:1806.09055 (2018).

[8]Cai, Han, Ligeng Zhu, and Song Han. “Proxylessnas: Direct neural architecture search on target task and hardware.” arXiv preprint arXiv:1812.00332 (2018).

本文转载自公众号才云 Caicloud（ID：Caicloud2015）。

原文链接：

https://mp.weixin.qq.com/s?__biz=MzIzMzExNDQ3MA==&mid=2650095901&idx=1&sn=3b65ffcd316493450eca182581c53e98&chksm=f08b7f14c7fcf602f0fa064c177ef2af7a84a530aa8a00140bfada49935562723aa94a39624c&scene=27#wechat_redirect