Architecture-aware Network Pruning for Vision Quality Applications

TL;DR

This paper introduces an iterative, architecture-aware pruning method for CNNs that reduces computational costs in vision quality tasks like low-light imaging and super-resolution without sacrificing quality.

Contribution

It presents a novel adaptive pruning algorithm that improves efficiency while maintaining image quality in high-resolution vision applications.

Findings

Reduced MAC by 58% in low-light imaging

Decreased MAC by 37% in super-resolution

Lowered memory bandwidth by 20-40%

Abstract

Convolutional neural network (CNN) delivers impressive achievements in computer vision and machine learning field. However, CNN incurs high computational complexity, especially for vision quality applications because of large image resolution. In this paper, we propose an iterative architecture-aware pruning algorithm with adaptive magnitude threshold while cooperating with quality-metric measurement simultaneously. We show the performance improvement applied on vision quality applications and provide comprehensive analysis with flexible pruning configuration. With the proposed method, the Multiply-Accumulate (MAC) of state-of-the-art low-light imaging (SID) and super-resolution (EDSR) are reduced by 58% and 37% without quality drop, respectively. The memory bandwidth (BW) requirements of convolutional layer can be also reduced by 20% to 40%.

Figures

Tables

Table 1. Table 1 : Detailed results. BW, considering only convolutional layers, consists of both weights and activations. Each weight and activation is represented with 4-byte floating-point numerical precision.

Network	Solution	$\#$ of MAC $(\times {10}^9)$		$\#$ of Weights $(\times {10}^3)$		$\#$ of Activations $(\times {10}^6)$		BW $($MByte/Inference$)$		Validation PSNR	Validation SSIM
SID	Original	560	(100%)	7757	(100%)	1915	(100%)	1922	(100%)	28.54	0.767
SID	Method-A	458	(82%)	6918	(89%)	1632	(85%)	1639	(85%)	28.54	0.768
SID	Method-B	354	(63%)	5275	(68%)	1485	(78%)	1491	(78%)	28.54	0.771
SID	Method-C	270	(48%)	5584	(72%)	1219	(64%)	1225	(64%)	28.54	0.769
SID	Method-D	236	(42%)	4241	(55%)	1169	(61%)	1173	(61%)	28.55	0.768
EDSR	Original	1428	(100%)	1367	(100%)	5076	(100%)	5077	(100%)	34.42	0.942
EDSR	Method-A	1085	(76%)	1037	(76%)	4481	(88%)	4481	(88%)	34.43	0.942
EDSR	Method-B	1085	(76%)	1037	(76%)	4481	(88%)	4481	(88%)	34.43	0.942
EDSR	Method-C	1085	(76%)	1037	(76%)	4481	(88%)	4481	(88%)	34.43	0.942
EDSR	Method-D	897	(63%)	857	(63%)	4083	(80%)	4083	(80%)	34.42	0.942

Equations

$$R_l=lo{g}_{10}(M_l/W_l)$$

$$T_l=T_b\times (1-S_l)\times R_l$$

Full text

**Index Terms— ** Pruning, Vision Quality, Network Architecture

1 Introduction

CNN is adopted as an essential ingredients in computer vision and machine learning areas [1, 2, 3]. Vision perception tasks including image classification, object detection and semantic segmentation are comprehensively investigated and associated with CNN. Even in image processing field such as super resolution, high dynamic range imaging and de-noising, CNN has progressive and promising improvement on image quality in recent years [4, 5].

However, compared to perception tasks, it requires higher computational complexity and BW requirements for vision quality tasks. MobileNetV1 [6] is designed with 569M MAC for ImageNet classification. On the other hand, in low-light photography SID [4] and super-resolution EDSR [5], it takes 560G MAC and 1.4T MAC per inference, respectively. It is more challenge to deploy CNN models on mobile devices for vision quality applications.

Network pruning [7] is an effective methodology toward performance optimization. Sparsity is defined as the ratio of the number of zero weights divided by the number of total weights. Better pruning algorithm delivers higher sparsity and reduces more MAC and BW correspondingly. However, quality drop is one of the major challenges in network pruning. Fig. 1 shows visible defect on SID even with only 0.1 PSNR degradation.

In this paper, we propose architecture-aware pruning to maximize sparsity and MAC reduction without quality-metric (PSNR or SSIM) drops. We also analyze the effects of MAC and BW reduction with different configurations associating with pruned structures. The proposed method focus on algorithms including but not limited to SID and EDSR.

2 Related Works

Network pruning has been widely explored in existing literatures. To answer which weight should be pruned, some works add evaluation functions to loss function, such as group lasso [8] and MAC regularization [9]. However, it is difficult to find a proper ratio between additional pruning-related loss and original loss. Others works create evaluation functions, including sensitivity [10, 11] and weight magnitude [12]. The sensitivity method computes the impact of weights on the training loss and removes low-impact weights. Weight magnitude method simply prune weight if its absolute value is less than the threshold, which is easier to be applied on large-scale CNNs. In this work, we use weight magnitude method to prune the network.

Pruning granularity. There are two granularity of pruning, fine-grained pruning [12, 13] and coarse-grained pruning [14, 15, 16, 17]. Fined-grained method prunes individual weights (i.e., within a filter kernel), whereas coarse-grained method extensively considers network structures (i.e., along the output and the input channels). According to [18], fine-grained pruning needs additional dedicated hardware to handle irregular sparsity. Coarse-grained method may obtain higher compression ratio without the need of compression header [19]. Therefore, we focus on coarse-grained output-channel-wise pruning.

Iterative pruning. To prevent catastrophic accuracy degradation, iterative pruning is viewed as an effective retraining procedure [20, 21, 22]. For vision quality applications, quality metrics are required as a reference judgement for termination of pruning procedure.

3 PROPOSED METHOD

3.1 Architecture-aware Pruning

An output channel is pruned if its maximum absolute weight value is less than magnitude threshold. For convolutional layer, the weight kernel has tensor shape $i\times o\times k\times k$, where $i$ is the number of input channels, $o$ is the number of output channels and $k$ is kernel size. Output-channel-wise pruning removes the weights along output channels. The kernel shape becomes $i\times(o-o^{\prime})\times k\times k$ if $o^{\prime}$ output channels are pruned. The output-channel pruned ratio is defined as $o^{\prime}/o$.

Once output channels of a layer are pruned, the corresponding input channels of the following layer are also removed. We defined one layer’s sparsity as $1-((i-i^{\prime})(o-o^{\prime}))/(i\times o)$, where $i^{\prime}$ is the number of pruned input channel in the layer. The network sparsity is defined as the ratio of the number of zero weights of a pruned network divided by the number of total weights of the original network.

3.1.1 Keep Layer Depth

Usually, in vision quality applications, each layer in network is semantically designed for quality-sensitive primitives, such as edge and chroma, with respect to different resolutions. Intensively removing a layer can severely degrade the quality. Therefore, we keep the network architecture by preserving minimum number of output channels.

3.1.2 Enhance MAC Efficiency

A pruned network with higher weight sparsity may not imply higher computation reduction. We define MAC/weight, $R_{l}$ (Eq. 1), for each layer as an indicator of MAC efficiency. $M_{l}$ and $W_{l}$ is number of MAC and weights of a layer, respectively. Fig. 2 shows that MAC/weight are much larger on both top and bottom layers on SID because of its U-Net [23] network topology. Therefore, to productively reduce computation, output channels of a layer with higher $R_{l}$ are tend to be pruned more because of higher magnitude threshold.

[TABLE]

3.1.3 Balance Pruned Output Channel

Residual block is universally used in network topology design such as in EDSR which is a variation of ResNet [24] with long shortcut. However, to prune output channels from residual block is arduous because of element-wise operations (i.e., element-wise ADD) or concatenation. Fig. 3 illustrates an example that a magnitude threshold is applied to 4-layer residual block. Because of element-wise ADD after layer Conv-B and layer Conv-D, the output channel of a given layer (Conv-D) and its preceding layer (Conv-B) with the same index (5) should be grouped and pruned at the same time. Conv-B layer and Conv-D layer have less pruned output channels compared to layer Conv-A and layer Conv-C.

We propose a guidance (Eq. 2) to prune output channels of residual block easier by increasing magnitude threshold on layers with lower ratio of pruned output channels. MAC efficiency mentioned in Sec. 3.1.2 is also applied. $S_{l}$ is the ratio of pruned output channels of layer $l$. $T_{b}$ is the original magnitude threshold base. Thus, output channels of a layer with lower $S_{l}$ have higher tendency to be pruned.

[TABLE]

3.2 Quality Metric Guarantee

To maintain quality metric (PSNR and SSIM) while maximizing pruned MAC, our algorithm prunes and retrains network iteratively. The iteration terminates when either the target quality-metric criteria or maximum training steps is reached. The proposed overall flow is shown in Algorithm 1.

4 Experimental Result

4.1 Experiment Setup

We generally investigate both SID for low-light photography and EDSR (baseline network, $\times$2) for super resolution. SID uses its own dataset [4] and EDSR adopts DIV2K dataset [25]. The input size of the network is set to the maximum image resolution in the datasets, 1424$\times$2128 for SID and 1020$\times$1020 for DIV2K, to calculate MAC and BW.

In SID dataset, we use images captured by Sony $\alpha$7SII camera as our training and validation data, which contains 280 pairs and 93 pairs, respectively. The pre-process stage is aligned with the setting in SID paper. In DIV2K dataset, we use the pre-process setting mentioned in [5] to generate 5,458,040 training patches from 800 training images and use 100 validation images as our validation data.

4.2 Result

The comprehensive analysis is elaborated in Table 1. We evaluate four distinct approaches. Method-A stands for magnitude threshold pruning without any structural hints. Method-B keeps the depth of network (Sec. 3.1.1). Method-C further considers MAC/weight ratio (Sec. 3.1.2) on the basis of Method-B. Method-D integrates all proposed techniques. All methods are conducted in company with quality-metric constraints (Sec. 3.2). For both SID and EDSR, it shows no PSNR and SSIM drop on all methods. Fig. 4 and Fig. 5 reveal the indistinguishable quality difference.

Keep Layer Depth. In SID, Method-B reduces MAC from 82% to 63% compared to Method-A. As shown in Fig. 6, Method-A may remove all the output channels of a layer due to not keeping layer depth, which leads to severe quality metric drops that cannot be recovered in retraining steps.

Enhance MAC Efficiency. In SID, Fig. 6 shows that Method-C prunes more weights on both top and bottom layers that have larger MAC/weight (Eq. 1). Therefore, Method-C reduces MAC from 63% to 48% compared to Method-B.

Balance Pruned Output Channel. Method-D increases 17% weight sparsity but only reduces 6% MAC compared to Method-C in SID. Fig. 6 illustrates that Method-D prunes less on top and bottom layers which have larger MAC/weight.

In EDSR, there is no difference among Method-A, Method-B and Method-C because no layer is pruned by Method-A and MAC/weight are identical for all layers (last layer could not be pruned) as shown in Fig. 2. Fig. 7 shows that Method-D reduces MAC from 76% to 63%, which is more than 10%, in shortcut-connected layers.

In summary, our methodology has significant reduction on both MAC and BW, which implies reduction on inference latency. For BW, we have 39% and 20% reduction on SID and EDSR, respectively. Our methodology also works well on complex network architecture.

5 Conclusion

To minimize computation complexity without quality drop on vision quality applications, our architecture-aware pruning is optimized for pruning more for complexity metric (e.g., MAC) on SID and shortcut-connected layers on EDSR. The MAC of SID and EDSR are reduced by 58% and 37%, respectively. Memory bandwidth is also reduced without degradation of PSNR, SSIM and subjective quality. The reduction of computation complexity and memory bandwidth could benefit on general mobile devices without special hardware design.