TorchX 与 Ax 大整合：更高效的多目标神经架构搜索

关注

原创 2023-02-27 01:04:19 1.3k 阅读 AI 写同款· GEO 优化 ›

TorchX一体化框架：从序列到结构的蛋白设计与工具对比计算蛋白设计正从单一工具链走向流程整合，其中结构预测、序列生成与功能评分是三个核心环节。传统工作流常因格式不统一、坐标不一致等问题在工具间搬运数据，耗费大量时间。TorchX作为一体化框架，通过统一数据协议将预测、设计、评分串成闭环，让抗体CDR优化、蛋白质-蛋白质相互作用改造等项目更高效。其设计模块利用结构预测的中间特征引导序列生成，评分模块提供可微反馈，提升候选质量。面对GenCDR、Boltz-2、OpenDDE、PPIFlow等专业工具，TorchX并不取代它们，而是充当工作流骨架，兼容外部工具。从阅读详情

NAS 的目的是发现神经网络的最佳架构。Torch 与 Ax 的整合为多目标探索神经架构的研究提供辅助。本文将展示用 Multi-objective Bayesian NAS 运行完全自动的神经架构搜索。

Ax 中的多目标优化 (Multi-Objective Optimization) 能够有效地探索神经架构搜索 (Neural Architecture Search) 中的 tradeoff ，比如模型性能和模型大小或延迟之间的 tradeoff。

这种方法已经在 Meta 公司成功应用于各种产品，如 On-Device AI。

本文提供了一个端到端的教程，帮助使用 TorchX。

Ax 简介

神经网络的规模和复杂性都在继续增长。开发最先进的架构通常是一个繁琐和耗时的过程，需要领域的专业知识和大量的工程努力。

为了克服这些挑战，人们提出了几种神经架构搜索 (NAS) 方法，以自动设计性能良好的架构，而不再需要人为监督 (human in the loop，HITL) 。

尽管样本效率很低，但像随机搜索和网格搜索这样的朴素法在超参数优化和 NAS 方面仍然很受欢迎。 NeurIPS 2019 和 ICLR 2020 进行的一项研究发现，80% 的 NeurIPS 论文和 88% 的 ICLR 论文使用 manual tuning 、随机搜索或网格搜索来调整 ML 模型超参数。

由于模型训练常常很耗时，并且可能需要大量的计算资源，所以尽量减少评估的 configurations 数量是非常重要的。

Ax 是一个用于黑箱优化的通用工具，允许用户使用贝叶斯优化等最先进的算法，在探索大型搜索空间时达到 sample efficient。

Meta 将 Ax 用于各种领域，包括超参数调整、NAS、用大规模 A/B 测试确定最佳产品设置、基础设施优化和设计尖端的 AR/VR 硬件领域。

许多 NAS 应用中，在多个 interest 的指标之间存在着自然的 tradeoff。例如，当在设备上部署模型时，会希望最大限度地提高模型性能（如准确率），同时尽量减少 competing metrics，如功耗、推理延迟或模型大小，以满足部署限制。

在许多情况下，模型性能的小幅下降换取大幅降低计算要求或预测的延迟是能够被接受的（在某些情况下能够同时提高准确率和降低延迟！）。

Sustainble AI 的关键推动力就是有效探索这种权衡蕴含的原则性方法。

Meta 利用 Ax 中的多目标贝叶斯 NAS (multi-objective Bayesian NAS) 成功地去探索这种权衡，并且该方法已经是用于优化 AR/VR 设备上的 ML 模型的例行用法。

除了 NAS 应用，Meta 还开发了可被用于 AR 光学系统的高维多目标优化的方法 MORBO。

Ax 中的 multi-objective Bayesian NAS

用 Ax 实现完全自动多目标 Nas

Ax 的 Scheduler 允许以闭环的方式异步运行实验，方法是不断将试验部署到外部系统，轮询结果，利用获取的数据生成更多的试验，并重复这一过程，直到满足停止条件。不需要人工干预或监督。这个 Scheduler 的特征有：

可定制并行性、故障容忍度和许多其他设置。
大量选择最先进的优化算法。
保存正在进行的实验（SQL DB 或 json），并从存储中恢复实验。
易于扩展到新的后端，以便远程运行试验评估。

下面是 Ax scheduler Tutorial 中的插图，总结了 Sceduler 如何与任何用于运行试验评估的外部系统互动。

Ax sceduler Tutorial

使用 Scheduler 运行自动 NAS，需要以下准备：

定义一个 Runner，负责发送一个具有特定架构的模型，在我们选择的平台上进行训练。（比如 Kubernetes，也可能只是本地上的 Docker 镜像）。下面的教程将用 TorchX 来处理训练作业的部署。
定义一个 Metric，负责从训练任务中获取客观指标（如准确度、模型大小、延迟）。下面的教程将用 Tensorboard 来记录数据，因此可以使用与 Ax 捆绑的 Tensorboard metric。

教程

教程展示如何使用 Ax 在流行的 MNIST 数据集上为一个简单的神经网络模型运行多目标 NAS。

虽然基本的方法论可用于更复杂的模型和更大的数据集，但这里选择了一个可在笔记本电脑上轻松运行的教程，可在一个小时内完成端到端的操作。

该例中，将会调整两个隐藏层的宽度、学习率、dropout probability、批尺寸和训练迭代的数量。目标是利用多目标贝叶斯优化来 trade off 性能（验证集的准确率）和模型大小（模型参数的数量）。

本教程使用了以下 PyTorch 库：

PyTorch Lightnig（用于指定模型和训练循环）

https://github.com/Lightning-AI/lightning

TorchX（用于远程/异步运行训练工作）

https://github.com/pytorch/torchx

BoTorch（为 Ax 的算法提供动力的贝叶斯优化库）

https://github.com/pytorch/botorch

完整可运行实例

结果

教程中进行的 NAS 优化的最终结果可以在下面的 tradeoff plot 中看到。这里，每个点都对应一个试验的结果，颜色代表其迭代数，星星表示对目标施加的阈值定义的参考点。

可以看到这种方法能够成功地探索验证精度和参数数量之间的 tradeoff，既能找到验证精度高的大模型，也能找到验证精度低的小模型。

根据性能要求和模型大小的限制，决策者现在可以选择使用或进一步分析哪个模型。

可视化

Ax 提供了一些可视化的东西帮助分析和理解实验结果。请重点关注对未知目标建模的高斯过程模型的性能，这些模型可以用来辅助更快地发现 promising configurations。

Ax 辅助更好地了解这些模型的准确率，以及它们在未见过的数据上的表现通过 leave-one-out cross-validation。

在下面的数字中，可以看出模型拟合的相当好--预测结果接近实际结果，预测的 95% 置信区间很好地覆盖了实际结果。

此外，能够看出模型大小 (num_params) 指标比验证精度 (val_acc）指标更容易建模。

对 Ax 总结

教程展示了如何使用 Ax 运行一个完全自动的多目标神经架构搜索。
使用 Ax 调度器，以完全异步的方式自动运行优化（这可以在本地完成），也可以通过远程部署试验到集群中（只需改变 TorchX 的 Scheduler 配置）。
Ax 中提供的最先进的多目标贝叶斯优化算法帮助有效地探索验证准确率和模型大小之间的 tradeoff。

高级功能

Ax 还有一些其他高级功能，上述教程中没有讨论。包括以下内容：

Early Stopping

在评估一个新的 candidate configutation 时，部分学习曲线通常在 NN 训练工作运行时就可以得到。

通过部分曲线中包含的信息来识别表现不佳的试验，以便及早停止，从而为更有希望的 promising candidate 释放计算资源。虽然在上述教程中没有演示，但 Ax 支持开箱即用的 Early Stopping。

https://ax.dev/versions/latest/tutorials/early_stopping/early_stopping.html

高维搜索空间

教程中使用了带有标准高斯过程的贝叶斯优化，以保持低运行时间。

然而，这些模型通常只能扩展到大约 10-20 个可调整的参数。新的 SAASBO 方法非常 sample-efficient 并且可以调整数百个参数。通过向 select_generation_strategy 传递 use_saasbo=True ，可以很容易地启用 SAASBO。

以上就是 TorchX 的介绍，之后 PyTorch 开发者社区将继续关注利用 Graph Transformations 优化生产型 PyTorch 模型的性能。

搜索 Hyperai01，备注「PyTorch」加入 PyTorch 技术交流群，了解 PyTorch 最新进展、最佳实践吧！

收录于