[Paper Review] aw_nas: A Modularized and Extensible NAS framework

Xuefei Ning, Changcheng Tang+8 more|arXiv (Cornell University)|Nov 25, 2020

Adversarial Robustness in Machine Learning22 references4 citations

TL;DR

aw_nas is a modular, extensible, and open-source Python framework for Neural Architecture Search (NAS) that supports diverse NAS algorithms and applications through well-defined, composable components. It enables efficient hardware-aware NAS by providing accurate latency and energy prediction models (e.g., LSTM-based) that reduce rMSE by up to 9.7× compared to naive summation, significantly improving cost estimation for CPU, GPU, and FPGA platforms.

ABSTRACT

Neural Architecture Search (NAS) has received extensive attention due to its capability to discover neural network architectures in an automated manner. aw_nas is an open-source Python framework implementing various NAS algorithms in a modularized manner. Currently, aw_nas can be used to reproduce the results of mainstream NAS algorithms of various types. Also, due to the modularized design, one can simply experiment with different NAS algorithms for various applications with awnas (e.g., classification, detection, text modeling, fault tolerance, adversarial robustness, hardware efficiency, and etc.). Codes and documentation are available at https://github.com/walkerning/aw_nas.

Motivation & Objective

• To provide a general, extensible, and easy-to-use NAS framework that supports diverse NAS algorithms and applications.
• To address the high computational cost of NAS by enabling efficient, modular component-based design for search space, controller, evaluator, and weights management.
• To support hardware-aware NAS by integrating accurate latency and energy prediction models across CPU, GPU, and FPGA platforms.
• To facilitate reproducibility and experimentation by offering open-source code and comprehensive documentation.
• To improve sample efficiency and scalability in NAS through modular, composable components and predictor-based search strategies.

Proposed method

• The framework uses a modular component architecture, with well-defined interfaces between components such as search space, controller, evaluator, weights manager, and trainer.
• It introduces a 'rollout' abstraction (BaseRollout) to standardize communication between components, supporting both discrete and differentiable architectures.
• The framework supports multiple NAS paradigms, including reinforcement learning, evolutionary search, differentiable NAS, and predictor-based methods.
• It implements parameter-sharing evaluation via a supernet and differentiable supernet, enabling fast architecture evaluation without retraining.
• It includes a hardware profiling pipeline that measures latency and energy on CPU, GPU, and FPGA platforms, and uses this data to train cost prediction models.
• Three types of cost prediction models are implemented: linear regression (1- and 2-variant), multilayer perceptron (MLP), and LSTM-based models that process block features sequentially.

Experimental results

Research questions

• RQ1 How can a NAS framework be designed to be modular, extensible, and reusable across diverse applications and NAS algorithms?
• RQ2 What is the impact of different cost prediction models on the accuracy of latency and energy estimation for neural networks on heterogeneous hardware platforms?
• RQ3 Can predictor-based NAS methods be effectively integrated into a modular framework to improve sample efficiency?
• RQ4 How does the performance of hardware cost prediction models compare to naive summation of primitive costs on platforms with complex memory and parallelism characteristics?
• RQ5 To what extent can the framework support hardware-aware NAS for real-world deployment in resource-constrained environments?

Key findings

• The LSTM-based cost prediction model achieved an rMSE of 0.174–0.187 for GPU latency estimation, representing a 3.6× improvement over naive summation (rMSE = 0.63).
• For FPGA energy estimation, the prediction models reduced rMSE from 4.94 (naive addition) to 0.51, achieving a 9.7× improvement.
• The framework's modular design enables seamless integration of diverse NAS algorithms, including reinforcement learning, evolutionary, differentiable, and predictor-based methods.
• Hardware profiling on CPU, GPU, and FPGA platforms was successfully automated, with cost prediction models trained on measured latency and energy of primitive blocks.
• The framework supports a wide range of applications, including image classification, object detection, NLP, fault tolerance, adversarial robustness, and hardware efficiency.
• The open-source release of cost prediction models and hardware assets (e.g., for Intel Xeon, RTX-2080Ti, and Xilinx ZCU102) enables reproducible and scalable hardware-aware NAS.