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: Working memory utilization for most CNN based models typically is more than what can be afforded on Tiny devices. Thus, it is often quite difficult to extract reasonable performance for visual tasks on tiny devices.  Authors attack this problem by a) introducing a patch based inference method that improves peak memory utilization at the cost of compute and b) introducing a redistribution process that reduces the overall compute while maintaining/improving performance.