Memory-efficient Patch-based Inference for Tiny Deep Learning

Tiny deep learning on microcontroller units ( MCUs ) is challenging due to the limited memory size . We find that the memory bottleneck is due to the imbalanced memory distribution in convolutional neural network ( CNN ) designs : the first several blocks have an order of magnitude larger memory usage than the rest of the network . To alleviate this issue , we propose a generic patch-by-patch inference scheduling , which operates only on a small spatial region of the feature map and significantly cuts down the peak memory . However , naive implementation brings overlapping patches and computation overhead . We further propose receptive field redistribution to shift the receptive field and FLOPs to the later stage and reduce the computation overhead . Manually redistributing the receptive field is difficult . We automate the process with neural architecture search to jointly optimize the neural architecture and inference scheduling , leading to MCUNetV2 . Patch-based inference effectively reduces the peak memory usage of existing networks by 4-8× . Co-designed with neural networks , MCUNetV2 sets a record ImageNet accuracy on MCU ( 71.8 % ) , and achieves > 90 % accuracy on the visual wake words dataset under only 32kB SRAM . MCUNetV2 also unblocks object detection on tiny devices , achieving 16.9 % higher mAP on Pascal VOC compared to the stateof-the-art result . Our study largely addressed the memory bottleneck in tinyML and paved the way for various vision applications beyond image classification . 1 Introduction . IoT devices based on tiny hardware like microcontroller units ( MCUs ) are ubiquitous nowadays . Deploying deep learning models on such tiny hardware will enable us to democratize artificial intelligence . However , tiny deep learning is fundamentally different from mobile deep learning due to the tight memory budget [ 27 ] : a common MCU usually has an SRAM smaller than 512kB , which is too small for deploying most off-the-shelf deep learning networks . Even for more powerful hardware like Raspberry Pi 4 , fitting inference into the L2 cache ( 1MB ) can significantly improve energy efficiency . These pose new challenges to efficient AI inference with a small peak memory usage . Existing work employs pruning [ 16 , 18 , 19 , 32 , 31 ] , quantization [ 16 , 54 , 53 , 37 , 46 , 9 , 39 ] , and neural architecture search [ 6 , 44 , 47 , 5 ] for efficient deep learning . However , these methods focus on reducing the number of parameters and FLOPs , but not the memory bottleneck . The tight memory budget limits the feature map/activation size , restricting us to use a small model capacity or a small input image size . Actually , the input resolutions used in existing tinyML work are usually small ( < 2242 ) [ 27 ] , which might be acceptable for image classification ( e.g. , ImageNet [ 11 ] , VWW [ 10 ] ) , but not for dense prediction tasks like objection detection : as in Figure 2 , the performance of object detection degrades much faster with input resolution than image classification . Such a restriction hinders the application of tiny deep learning on many real-life tasks ( e.g. , person detection ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 4x4patch ( int8 ) A cc ur ac y/ m A P ( % ) 60 65 70 75 Image Resolution 160 224 288 352 VOC mAP ImgNet Top-1 larger degrade VOC mAP ImgNet Top-1 larger degrade Image Resolution A cc ur ac y/ m A P ( % ) 160 224 228 352 75 70 65 60 Figure 2 . Detection is more sensitive to smaller resolutions . We perform an in-depth analysis on memory usage of each layer in efficient network designs and find that they have a very imbalanced activation memory distribution . Take MobileNetV2 [ 41 ] as an example , as shown in Figure 1 , only the first 5 blocks have a high peak memory ( > 450kB ) , becoming the memory bottleneck of the entire network . The remaining 13 blocks have a low memory usage , which can easily fit a 256kB MCU . The peak memory of the initial memory-intensive stage is 8× higher than the rest of the network . Such a memory pattern leaves a huge room for optimization : if we can find a way to “ bypass ” the memory-intensive stage , we can reduce the overall peak memory by 8× . In this paper , we propose MCUNetV2 to address the challenge . We first propose a patch-by-patch execution order for the initial memory-intensive stage of CNNs ( Figure 3 ) . Unlike conventional layer-by-layer execution , it operates on a small spatial region of the feature map at a time , instead of the whole activation . Since we only need to store the feature of a small patch , we can significantly cut down the peak memory of the initial stage ( blue to yellow in Figure 3 ) , allowing us to fit a larger input resolution . However , the reduced peak memory comes at the price of computation overhead : in order to compute the non-overlapping output patches , the input image patches need to be overlapped ( Figure 3 ( b ) ) , leading to repeated computation . The overhead is positively related to the receptive field of the initial stage : the larger the receptive field , the larger the input patches , which leads to more overlapping . We further propose receptive field redistribution to shift the receptive field and workload to the later stage of the network . This reduces the patch size as well as the computation overhead caused by overlapping , without hurting the performance of the network . Finally , patch-based inference brings a larger design space for the neural network , giving us more freedom trading-off input resolution , model size , etc . We also need to minimize the computation overhead under patch-based execution . To explore such a large and entangled space , we propose to jointly design the optimal deep model and its inference schedule with neural architecture search given a specific dataset and hardware . Patch-based inference effectively reduces the peak memory usage of existing networks by 4-8× ( Figure 5 ) . The results are further improved when co-designing neural architecture with inference scheduling . On ImageNet [ 11 ] , we achieve a record accuracy of 71.8 % on MCU ( Table 2 ) ; on visual wake words dataset [ 10 ] , we are able to achieve > 90 % accuracy under only 32kB SRAM , which is 4.0× smaller compared to MCUNetV1 [ 27 ] , greatly lowering the boundary of tiny deep learning ( Figure 7 ) . MCUNetV2 further unlocks the possibility to perform dense prediction tasks on MCUs ( e.g. , object detection ) , which was not practical due to the limited input resolution . We are able to achieve 64.6 % mAP under 256kB SRAM constraints and 68.3 % under 512kB , which is 16.9 % higher compared to the existing state-of-the-art solution , making object detection applicable on a tiny ARM Cortex-M4 MCU . Our contributions can be summarized as follows : • We systematically analyze the memory usage pattern of efficient CNN designs and find that they suffer from a imbalanced memory distribution , leaving a huge room for optimization . • We propose a patch-based inference scheduling to significantly reduce the peak memory required for running CNN models , together with receptive field redistribution to minimize the computation overhead . • With the joint design of network architecture and inference scheduling , we achieve a record performance for tiny image classification and objection detection on MCUs . Our work largely addressed the memory bottleneck for tinyML , paving the way for various vision applications . 2 Understanding the Memory Bottleneck of Tiny Deep Learning . We systematically analyze the memory bottleneck of CNN models . Imbalanced memory distribution . As an example , we provide the per-block peak memory usage of MobileNetV2 [ 41 ] in Figure 1 . The profiling is done in int8 ( details in Section 4 ) . We can observe a clear pattern of imbalanced memory usage distribution . The first 5 blocks have large peak memory , exceeding the memory constraints of MCUs , while the remaining 13 blocks easily fit 256kB memory constraints . The third block has 8× larger memory usage than the rest of the network , becoming the memory bottleneck . We also inspect other efficient network designs and find the phenomenon quite common across different CNN backbones , even for models specialized for memory-limited microcontrollers [ 27 ] . The detailed statistics are provided in the supplementary . We find that this situation applies to most single-branch or residual CNN designs due to the hierarchical structure * : after each stage , the image resolution is down-sampled by half , leading to 4× fewer pixels , while the channel number increases only by 2× [ 42 , 17 , 22 ] or by an even smaller ratio [ 41 , 21 , 45 ] , resulting in a decreasing activation size . Therefore , the memory bottleneck tends to appear at the early stage of the network , after which the peak memory usage is much smaller . Challenges and opportunities . The imbalanced memory distribution significantly limits the model capacity and input resolution executable on MCUs . In order to accommodate the initial memoryintensive stage , the whole network needs to be scaled down even though the majority of the network already has a small memory usage . It also makes resolution-sensitive tasks ( e.g. , object detection ) difficult , as a high-resolution input will lead to large initial peak memory . Consider the first convolutional layer in MobileNetV2 [ 41 ] with input channels 3 , output channels 32 , and stride 2 , running it on an image of resolution 224 × 224 requires a memory of 3 × 2242 + 32 × 1122 = 539kB even when quantized in int8 , which can not be fitted into microcontrollers . On the other hand , if we can find a way to “ bypass ” the initial memory-intensive stage , we can greatly reduce the peak memory of the whole network , leaving us a large room for optimization . 3 MCUNetV2 : Memory-Efficient Patch-based Inference . 3.1 Breaking the Memory Bottleneck with Patch-based Inference . We propose to break the memory bottleneck of the initial layers with patch-based inference ( Figure 3 ) . Existing deep learning inference frameworks ( e.g. , TensorFlow Lite Micro [ 1 ] , TinyEngine [ 27 ] , microTVM [ 8 ] , etc . ) use a layer-by-layer execution . For each convolutional layer , the inference library first allocates the input and output activation buffer in SRAM , and releases the input buffer after the whole layer computation is finished . Such an implementation makes inference optimization easy ( e.g. , im2col , tiling , etc . ) , but the SRAM has to hold the entire input and output activation for * some CNN designs have highly complicated branching structure ( e.g. , NASNet [ 56 ] ) , but they are generally less efficient for inference [ 34 , 44 , 6 ] ; thus not widely used for edge computing . each layer , which is prohibitively large for the initial stage of the CNN . Our patch-based inference runs the initial memory-intensive stage in a patch-by-patch manner . For each time , we only run the model on a small spatial region ( > 10× smaller than the whole area ) , which effectively cuts down the peak memory usage . After this stage is finished , the rest of the network with a small peak memory is executed in a normal layer-by-layer manner ( upper notations in Figure 1 ) . We show an example of two convolutional layers ( with stride 1 and 2 ) in Figure 3 . For conventional per-layer computation , the first convolutional layer has large input and output activation size , leading to a high peak memory . With spatial partial computation , we allocate the buffer for the final output and compute its values patch-by-patch . In this manner , we only need to store the activation from one patch instead of the whole feature map . Note that the first activation is the input image , which can be partially decoded from a compressed format like JPEG and does not require full storage . Computation overhead . The significant memory saving comes at the cost of computation overhead . To maintain the same output results as per-layer inference , the non-overlapping output patches correspond to overlapping patches in the input image ( the shadow area in Figure 3 ( b ) ) . This is because convolutional filters with kernel size > 1 contribute to increasing receptive fields . The bordering pixel on the output patches is dependent on the inputs from neighboring patches . Such repeated computation can increase the overall network computation by 10-17 % even under optimal hyper-parameter choice ( Figure 5 ) , which is undesirable for low-power edge devices .

Summary: The main problem the paper tries to tackle is reducing the peak memory footprint of the neural networks by deviating from the usual layer-by-layer execution. The main intuition seems to be very similar to tiling combined with layer fusion-like ideas. However, the main spark of the paper seems to be leveraging this dimension in the context of Neural Architecture Search, which seems to be novel.