FBNetV2: Differentiable Neural Architecture Search for Spatial and Channel Dimensions

Introduction

Deep neural networks have led to significant progress in many research areas and applications, such as computer vision and autonomous driving. Despite this, designing an efficient network for resource-constrained settings remains a challenging problem. Initial directions involved compressing existing networks or building small networks . However, the design space can easily contain more than $10^{18}$ candidate architectures , making manual design choices sub-optimal and difficult to scale. In lieu of manual tuning, recent work uses neural architecture search (NAS) to design networks automatically.

Previous NAS methods utilize reinforcement learning (RL) techniques or evolutionary algorithms (EAs). However, both methods are computationally expensive and consume thousands of GPU hours . As a result, recent NAS literature focuses on differentiable neural architecture search (DNAS); DNAS searches over a supergraph that encompasses all candidate architectures, selecting a single path as the final neural network. Unlike conventional NAS, DNAS can search large combinatorial spaces in the time it takes to train a single model . One class of DNAS methods, based on DARTS , suffer from two significant limitations :

Memory costs bound the search space. Short of paging in and out tensors, the supergraph and feature maps must reside in GPU memory for training, which limits the search space.

Cost grows linearly with the number of options per layer. This means that each new search dimension introduces combinatorially more options and combinatorial memory and computational costs.

The other class of DNAS methods, not based on DARTS, suffer from similar issues: For example, ProxylessNAS tackles the memory constraint by training only one path in the supergraph each iteration. However, this means ProxylessNAS would take a prohibitively long time to converge on an order-of-magnitude larger search space. These memory and computation issues, for all DNAS methods, prevent us from expanding the search space to explore larger spaces of configurations. Noting that feature maps typically dominate memory cost , we propose a formulation of DNAS (Fig. 1) called DMaskingNAS (Fig. 2) that increases the search space size by orders of magnitude. To accomplish this, we represent multiple channel and input resolution options in the supergraph with masks, which carry negligible memory and computational costs. Furthermore, we reuse feature maps for all options in the supergraph, which enables nearly constant memory cost with increasing search space sizes. These optimizations yield the following three contributions:

A memory and computationally efficient DNAS that optimizes both macro- (resolution, channels) and micro- (building blocks) architectures jointly in a $10^{14}\times$ larger search space using differentiable search. To the best of our knowledge, we are the first to tackle this problem using a differentiable search framework supergraph, with substantially less computational cost and roughly constant memory cost.

A masking mechanism and effective shape propagation for feature map reuse. This is applied to both the spatial and channel dimensions in DNAS.

State-of-the-art results on ImageNet classification. With only 27 hours on 8 GPUs, our searched compact models lead to substantial per-parameter, per-FLOP accuracy improvements. The searched models outperform all previous state-of-the-art neural networks, both manually and automatically designed, small and large.

Related Work

Hand-crafted, efficient neural networks see two predominant approaches: (1) compressing existing architectures and (2) designing compact architectures from scratch.

Network compression includes both architectural and non-architectural modifications. One non-architectural approach is low-bit quantization, where weights and activations alike may be represented with fewer bits. For example, Wang et al. propose hardware-aware automated quantization, which achieves a $1.4$-$1.95\times$ latency reduction on MobileNet . These techniques are orthogonal to and can be combined with the methods in this paper. Alternatively, architectural modifications include network pruning , where various heuristics govern layer-wise or channel-wise pruning. For example, Han et al. show that magnitude-based pruning can reduce parameter count by orders of magnitude without accuracy loss, and NetAdapt utilizes a filter pruning algorithm that achieves a 1.2$\times$ speedup for MobileNetV2. However, with heuristics-based simplifications, pruning methods train potential architectures separately, one after another – in some cases, pruning methods consider only one architecture .

Compact architecture design aims to directly construct efficient networks, rather than trim an expensive one . For example, MobileNet and MobileNetV2 achieve substantial efficiency improvements by exploiting a depth-wise convolution and an inverted residual block, respectively. ShuffleNetV2 shrinks the model size utilizing low-cost group convolutions. Tan et al. propose a compound scaling method, obtaining a family of architectures that achieve state-of-the-art accuracy with an order of magnitude fewer parameters than previous convolutional networks . However, these models rely on finely-tuned, manual decisions that are bested by automatic design.

Neural architecture search automates the design of state-of-the-art neural networks. Zoph et al. first proposed using RL for automated neural network design in . This and other early NAS approaches are based on RL and EA . However, both approaches consume substantial computational resources.

Later works utilize various techniques to reduce the computational cost of search. One such technique formulates the architecture search problem as a path-finding process in a supergraph . Among them, gradient-based NAS has emerged as a promising tool. Wu et al. show that gradient-based, differentiable NAS yields state-of-the-art compact architectures with $421\times$ less search cost than RL-based approaches. Another direction is to exploit a performance predictor to guide the search process . Such approaches explore the search space by trimming progressively and lead to significant reductions in search cost.

Stamoulis et al. introduce weight-sharing to further reduce the computational cost of search. However, kernel weight-sharing doesn’t address the primary drawback of DARTS, namely a memory bottleneck yielding small search space size: Say a “mixed kernel” contains weights shared between a $3\times 3$ and $5\times 5$. Since it is impossible to extract a $3\times 3$ convolution’s outputs from a $5\times 5$’s (and vice versa), this mixed kernel still convolves $2\times$ and still stores 2 feature maps for backpropagation. Thus, 2 kernel-weight-sharing convolutions induce memory and computational costs of 2 vanilla convolutions.

Searching along spatial and channel dimensions has been studied both with and without NAS. Liu et al develop a NAS variant that searches over varying strides for semantic segmentation. However, this method suffers from increasing memory cost as the number of possible input resolutions grows. As described above, network pruning suffers from inefficient and sequential exploration of architectures, one-by-one. Yu et al amend this partially by creating a batchnorm invariant to the number input channels; after training the “supergraph” they see competitive accuracy without further training, for each possible subset of channels. Yu et al expand on these slimmable networks by introducing a test-time greedy channel selection procedure. However, these methods are orthogonal to and can be combined with DMaskingNAS, as we train the sampled architecture from scratch. To address these concerns, our algorithm jointly optimizes over multiple input resolutions and channel options simultaneously, increasing memory cost only negligibly as the number of options grows. This allows DMaskingNAS to support orders of magnitude more possible architectures, under existing memory constraints.

Method

We propose DMaskingNAS to search over spatial and channel dimensions, summarized in Fig. 2. The search space would be computationally prohibitive and ill-formed without the optimizations described below; our approach makes it possible to search this expanded search space (Table 1) over channels and input resolutions.

To support searches over varying numbers of channels, previous DNAS methods simply instantiate a block for every channel option in the supergraph. For a convolution with $k$ filters, this could mean up to $k(k+1)/2\sim O(k^{2})$ convolutions. Previous channel pruning methods suffer from a similar drawback: each option must be trained separately, finding the “optimal” channel count in one shot or iteratively. Furthermore, even without saturating the maximum number of possibilities, there are two problems, the first of which makes this search impossible:

Incompatible dimensions: DNAS is divided into several “cells”. In each cell, we consider a number of different block options; the outputs of all options are combined in a weighted sum. This means that all block outputs must align dimensions. If each block adopts convolutions with different number of filters, each output will have a different number of channels. As a result, DNAS could not perform a weighted sum.

Slower training, increased memory cost: Even with a workaround, with this naïve instantiation, each convolution with a different channel option must be run separately, resulting in a $O(k)$ increase in FLOP cost. Furthermore, each output feature map must be stored separately in memory.

To address the aforementioned issues, we handle the incompatibility (Fig. 3, Step A): consider a block $b$ with varying numbers of filters, where $b_{i}$ denotes this block with $i$ filters. The maximum number of filters is $k$. The outputs of all blocks are then zero-padded to have $k$ channels (Fig. 3, Step B). Given input $x$, the Gumbel Softmax output is thus the following, with Gumbel weights $g_{i}$:

Finally, with this approximation, we can handle the computational complexity of the naïve channel search approach: this is equivalent to computing the aggregate mask and running the block $b$ only once (Fig. 3, Step E).

This approximation only requires one forward pass and one feature map, inducing no additional FLOP or memory costs other than the negligible $M$ term in Eq. 3 (Fig. 2, Channel Masking). Furthermore, the approximation falls short of equivalence only because weights are shared, which is shown to reduce train time and boost accuracy in DNAS . This allows us to search the number of output channels for any block, including related architectural decisions such as the expansion rate in an inverted residual block.

2 Input Resolution Search

For spatial dimensions, we search over input resolutions. As with channels, previous DNAS methods would simply instantiate each block with every input resolution. This naïve method’s downfalls are twofold: increased memory cost and incompatible dimensions. As before, we address both issues directly by zero-padding the result. However, there are two caveats:

Pixel misalignment: means padding cannot occur naïvely as before. It would not make sense to zero-pad the periphery of the image, since the sum in Eq. 1 would result in misaligned pixels (Fig. 4, B). To handle pixel misalignment, we zero-pad such that zeros are interspersed spatially (Fig. 4, C). This zero-padding pattern is uniform; except for the zeros, this is a nearest neighbors upsampling. For example, a $2\times$ increase in size would involve zero-padding every other row and column. Zero-padding instead of upsampling minimizes “pixel contamination” across input resolutions (Fig. 5).

Receptive field misalignment: Since subsets of the feature map correspond to different resolutions, naïvely convolving over the full feature map would result in a reduced receptive field (Fig. 4, D). To handle receptive field misalignment, we convolve over subsampled input instead. (Fig. 4, E). Using Gumbel Softmax, we arrive at “resolution subsampling” in Fig. 2.

NASNet introduces a similar notion of combining hidden states. These combinations are also used to efficiently explore a combinatorially large search space but are used to determine – instead of input resolution or channels – the number of times to repeat a searched cell. With the above insights, the input resolution search thus incurs constant memory cost, regardless of the number of input resolutions. On the other hand, computational cost increases sub-linearly as the number of resolutions grows.

3 Effective Shape Propagation

Note this calculation for effective shape is only used during training. In our formulation of the weighted sum Eq. 1, the output $y$ retains the maximum number of channels. However, there exists a non-integral number of effective channels: say a 16-channel output has Gumbel weight $g_{i}=0.8$ and a 12-channel output has weight $g_{i}=0.2$. This means the effective number of channels is $0.8*16+0.2*12=15.2$. These effective channels are necessary for both FLOP and parameter computation, as assigning higher weight to more channels should incur a larger cost penalty. This effective shape is how we realize effective resource costs introduced in previous works : First, define the gumbel softmax weights as

with sampling parameter $\alpha$, Gumbel noise $\epsilon$, temperature $\tau$. For a convolution with Gumbel Softmax in the $l^{th}$ layer, we define its effective output shape $\bar{S}^{l}_{out}$ in Eq. 7 using effective output channel ($\bar{C}^{l}_{out}$, Eq. 5), and effective height, width ($\bar{h}^{l}_{out},\bar{w}^{l}_{out}$, Eq. 6).

with batch size $n$, effective input width $\bar{w}_{in}$ and height $\bar{h}_{in}$.

For a convolution layer without a Gumbel Softmax, effective output shape simplifies to Eq. 8, where effective channel count is equal to actual channel count. For a depth-wise convolution, effective output shape simplifies to Eq. 9, where effective channel count is simply propagated.

with actual output channel count $C_{out}$, effective input channel count $\bar{C}_{in}$. Then, we define the cost function for the $l^{th}$ layer as follow:

with $\gamma$ convolution groups. The effective input channels for the $(l+1)^{th}$ layer are $\bar{C}^{l+1}_{in}=\bar{C}^{l}_{out}$. The total training loss consists of (1) cross-entropy loss and (2) total cost, which is the sum of cost from all layers: $\text{cost}_{total}=\Sigma_{l}\text{cost}^{l}$.

In the forward pass, for all convolutions, we calculate and return both the output tensor and effective output shape. Additionally, $\tau$ in the Gumbel Softmax Eq. 4 decreases throughout training, , forcing $g^{l}$ to approach a one-hot distribution. $\text{argmax}_{i}g^{l}_{i}$ would thus select a path of blocks in the supergraph; a single channel and expansion rate option for each block; and a single input resolution for the entire network. This final architecture is then trained. Note this final model does not employ masking or require effective shapes.

Experiments

We use DMaskingNAS to search for convolutional network architectures under different objectives. We compare our search space, performance of searched models, and search cost to previously state-of-the-art networks. Detailed numerical results are listed in Table 4.

We implement DMaskingNAS using PyTorch on 8 Tesla V100 GPUs with 16GB memory. We use DMaskingNAS to search for convolutional neural networks on the ImageNet (ILSVRC 2012) classification dataset , a widely-used NAS evaluation benchmark. We use the same training settings as reported in : we randomly select 10% of classes from the original 1000 classes and train the supergraph for 90 epochs. In each epoch, we train the network weights with 80% of training samples using SGD. We then train the Gumbel Softmax sampling parameter $\alpha$ with the remaining 20% using Adam . We set initial temperature $\tau$ to 5.0 and exponentially anneal by $e^{-0.045}\approx 0.956$ every epoch.

2 Search Space

Previous cell-level searches produced fragmented, complicated, and latency-unfriendly blocks. Thus, we adopt a layer-wise search space for known, latency-friendly blocks.

Table 3 describes the micro-architecture search space: the block structure is inspired by and sequentially consists of a $1\times 1$ point-wise convolution, a $3\times 3$ or $5\times 5$ depth-wise convolution, and another $1\times 1$ point-wise convolution. Table 2 describes the macro-architecture. The search space contains more than $10^{35}$ candidate architectures, which is $10^{14}\times$ larger than DNAS’s .

3 Memory Cost

Our memory optimizations yield a $\sim$1MB increase in memory cost for every 2 orders of magnitude the channel search space grows by; for context, this 1 MB increase is just 0.1% of the total memory cost during training. This is due to our feature map reuse as described in Sec. 3.1. We compare memory costs for DNAS and DMaskingNAS as the number of channel options increases (Fig. 7, left). With only 8 channel options for each convolution, DNAS fails to fit in memory during training, exceeding the 16GB memory supported by a Tesla V100 GPU. On the other hand, DMaskingNAS supports 32-option channel search, for a $32^{22}\sim 10^{33}$ in search space size (given our 22-layer search space), at nearly constant memory cost. Here, $k$-option channel search means that for each convolution with $c$ channels, we search over $\{c/k,2c/k,...,c\}$ channels. To compare larger numbers of channel options, we reduce the number of blocks options in the search space (Fig. 7, right). To compute memory cost, we average the maximum memory allocated during each training step, across 10 epochs.

4 Search for ImageNet Models

FLOP-efficient models: We first use DMaskingNAS to find compact models (Fig. 6) for low computational budgets, with models ranging from 50 MFLOPs to 300 MFLOPs in Fig. 8. The searched FBNetV2s outperform all existing networks.

Storage-efficient models: Many real world scenarios face limited on-device storage space. Thus, we next perform searches for models minimizing parameter count, in Fig. 9. With similar or smaller model size (4M parameters), FBNetV2 achieves 2.6% and 2.9% absolute accuracy gains over MobileNetV3 and FBNet , respectively.

Large models: We finally use DMaskingNAS to explore larger models for high-end devices. We compare FBNetV2-Large with networks of 300+ MFLOPs in Fig. 10.

Conclusions

We propose a memory-efficient algorithm, drastically expanding the search space for DNAS by supporting searches over spatial and channel dimensions. These contributions target the main bottleneck for DNAS – high memory cost that induces constraints on the search space size – and yield state-of-the-art performance.

Acknowledgements In addition to NSF CISE Expeditions Award CCF-1730628, UC Berkeley research is supported by gifts from Alibaba, Amazon Web Services, Ant Financial, CapitalOne, Ericsson, Facebook, Futurewei, Google, Intel, Microsoft, Nvidia, Scotiabank, Splunk and VMware. This material is based upon work supported by the National Science Foundation Graduate Research Fellowship under Grant No. DGE 1752814.

References