SurgeNAS: a comprehensive surgery on hardware-aware differentiable neural architecture search

Abstract

Differentiable neural architecture search (NAS) is an emerging paradigm to automate the design of top-performing convolutional neural networks (CNNs). However, previous differentiable NAS methods suffer from several crucial weaknesses, such as inaccurate gradient estimation, high memory consumption, search fairness, etc. More importantly, previous differentiable NAS works are mostly hardware-agnostic since they only search for CNNs in terms of accuracy, ignoring other critical performance metrics like latency. In this work, we introduce a novel hardware-aware differentiable NAS framework, namely SurgeNAS, in which we leverage the one-level optimization to avoid inaccuracy in gradient estimation. To this end, we propose an effective identity mapping regularization to alleviate the over-selecting issue. Besides, to mitigate the memory bottleneck, we propose an ordered differentiable sampling approach, which significantly reduces the search memory consumption to the single-path level, thereby allowing to directly search on target tasks instead of small proxy tasks. Meanwhile, it guarantees the strict search fairness. Moreover, we introduce a graph neural networks (GNNs) based predictor to approximate the on-device latency, which is further integrated into SurgeNAS to enable the latency-aware architecture search. Finally, we analyze the resource underutilization issue, in which we propose to scale up the searched SurgeNets within \textit{Comfort Zone} to balance the computation and memory access, which brings considerable accuracy improvement without deteriorating the execution efficiency. Extensive experiments are conducted on ImageNet with diverse hardware platforms, which clearly show the effectiveness of SurgeNAS in terms of accuracy, latency, and search efficiency.