Semi-Supervised Semantic Segmentation via Adaptive Equalization Learning

1 Introduction . Supervised semantic segmentation requires pixel-level labeling , which is expensive and timeconsuming . This paper is interested in semi-supervised semantic segmentation , which can greatly reduce the efforts of pixel-level annotation , yet may maintain reasonably high accuracy . One problem of common semantic segmentation datasets is that the pixel categories tend to be imbalanced , e.g. , the pixel amount of head classes can be hundreds of times larger than that of tailed classes in the widely used Cityscapes dataset [ 1 ] . The situation is more serious in the semi-supervised setting where tailed classes may have extremely few samples . We note that recent approaches are mainly dedicated to the design of consistency regularization [ 2 , 3 , 4 , 5 , 6 , 7 ] and pseudo-labeling [ 8 ] , almost all of which neglect the imbalance problem and treat each category equally , leading to a biased training . These approaches may even harm the learning of tailed classes , as inaccurate predictions or pseudo labels of under-performing categories could falsely guide the learning on unlabeled data . * This work is done when Hanzhe Hu was an intern in MSRA . † Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . This paper aims to alleviate this biased training problem . We propose a novel Adaptive Equalization Learning ( AEL ) framework , which adaptively balance the training of different categories as shown in Figure 1 . Our design follows two main principles : 1 ) increasing the proportion of training samples from the under-performing categories ; 2 ) tilting training towards under-performing categories . Concretely , we maintain a confidence bank to dynamically record the category-wise performance at each training step , which indicates the current performance of each category . Following principle 1 ) , we propose two data augmentation approaches named adaptive Copy-Paste and adaptive CutMix , which give more chance for under-performing categories to be copied or cut . Following principle 2 ) , we present an adaptive equalization sampling strategy to encourage pixels from under-performing categories to be sufficiently trained . In addition , we also introduce a simple yet effective re-weighting strategy which takes the model predictions into account to alleviate the issue that semi-supervised learning usually suffers from the training noise . Experimentally , by using the DeepLabv3+ with ResNet-101 backbone , the proposed AEL outperforms state-of-the-art methods by a large margin on the Cityscapes and PASCAL VOC 2012 benchmarks under various data partition protocols . Specifically , it achieves 74.28 % , 75.83 % and 77.90 % on Cityscapes dataset under 1/32 , 1/16 and 1/8 protocols , which is +16.39 % , +12.87 % and +8.09 % better than the supervised baseline . When evaluated on PASCAL VOC 2012 benchmark , it achieves 76.97 % , 77.20 % and 77.57 % under 1/32 , 1/16 and 1/8 protocols , which is +6.83 % , +6.60 % and +4.45 % better than the supervised baseline . Moreover , the proposed approach also proves to improve the segmentation model trained on the full Cityscapes train set by +1.03 % by leveraging 5 , 000 images from the Cityscapes coarse set as unlabeled data , achieving 81.95 % . 2 Related Work . Semi-Supervised Learning . Recent years have witnessed a significant progress in the SSL field . Most of them can be categorized into consistency regularization , entropy minimization [ 9 ] and pseudo-labeling . Consistency regularization [ 10 , 11 , 12 ] enforces consistency in predictions between different views of unlabeled data . Pseudo-labeling [ 13 , 14 ] trains the model on the unlabeled data with pseudo labels generated from the model ’ s own predictions . Furthermore , [ 10 , 15 , 16 , 17 , 18 ] use a low softmax temperature to sharpen the predictions of unlabeled set . Our method refers to Mean Teacher [ 11 ] and FixMatch [ 14 ] when designing our basic framework . Semi-Supervised Semantic Segmentation . Existing semi-supervised semantic segmentation methods mainly focus on the design of consistency regularization and pseudo-labeling . Cutmix-Seg [ 2 ] applies CutMix augmentation on the unlabeled data . CCT [ 4 ] introduces a feature-level perturbation and enforces consistency among the predictions of different decoders . GCT [ 19 ] performs network perturbation by using two differently initialized segmentation models and encourages consistency between the predictions from the two models . PseudoSeg [ 8 ] focuses on improving the quality of pseudo labels . Though achieving satisfactory improvements over the supervised baseline , none of the aforementioned methods explore the biased learning issue in semi-supervised semantic segmentation . Class Imbalance in Semi-Supervised Learning . Although SSL has been extensively studied , class imbalance problem in SSL is relatively under-explored , especially for semantic segmentation . Yang et al . [ 20 ] demonstrate that leveraging unlabeled data can alleviate imbalance issue . Hyun et al . [ 21 ] propose a suppressed consistency loss for class-imbalanced image classification problems . CReST [ 22 ] introduces a self-training framework for imbalanced SSL . Our method , though not explicitly targeting at the class imbalance problem , focuses on improving the performance of underperforming categories which are mostly tailed classes . Moreover , we refer to the ideas of resampling [ 23 , 24 ] and re-weighting [ 25 , 26 ] , which are designed for class imbalance problem . 3 Method . Given a labeled set Dl = { ( xl i , yl i ) } and an unlabeled set Du = { xu i } , the objective of semisupervised semantic segmentation is to learn a segmentation model by efficiently leveraging both labeled and unlabeled data . In this section , we first present an overview of the proposed AEL in Section 3.1 . Then we describe our basic framework for semi-supervised semantic segmentation in Section 3.2 . Finally , the details of AEL are introduced in Section 3.3 . 3.1 Overview . Figure 2 displays an overview of AEL , which is a data-efficient framework for semi-supervised semantic segmentation . It is composed of two parts : 1 ) a basic framework which contains a teacher model for pseudo-labeling and a student model for online learning ; 2 ) dedicated modules which encourages the under-performing categories to be sufficiently trained by effectively leveraging both labeled and unlabeled data . We use the proposed confidence bank to dynamically record the categorywise performance during training , and thus we can easily identify which categories are not sufficiently trained . For those unsatisfactory categories , we present two data augmentation methods to increase their frequency of occurrence in a training batch , namely adaptive CutMix which is applied on the unlabeled data , and adaptive Copy-Paste which is applied on the labeled data . To make the model towards the unbiased learning , we propose the adaptive equalization sampling and dynamic re-weighting strategies to involve enough samples from the under-performing categories into the training , and alleviate the noise raised by pseudo-labeling simultaneously . 3.2 Basic Framework . We first set up a basic framework for semi-supervised semantic segmentation . The framework consists of a student model and a teacher model . The teacher model has the same architecture as the student model , but uses a different set of weights which are updated by exponential moving average ( EMA ) of the student model [ 11 ] . Following FixMatch [ 14 ] , we use the teacher model to generate a set of pseudo labels Ŷ = { ŷi } on the weakly augmented unlabeled data Du . Subsequently , the student model is trained on both labeled data Dl ( of weak augmentation ) with the ground-truth and unlabeled data Du ( of strong augmentation ) with the generated pseudo labels Ŷ . We use standard random resize and random horizontal flip as the weak augmentation . Strong augmentation includes CutMix [ 27 ] and all data augmentation strategies used in the weak augmentation . The overall loss consists of the supervised loss Ls and the unsupervised loss Lu : Ls = 1 Nl NlX i=1 1 WH WHX j=1 ` ce ( yij , pij ) , ( 1 ) Lu = 1 Nu NuX i=1 1 WH WHX j=1 ` ce ( ŷij , pij ) , ( 2 ) where pij is the prediction of the j-th pixel in the i-th labeled ( or unlabeled ) image , Nl and Nu denote the number of labeled images and unlabeled images in a training batch , W and H represent the width and height of the input image , and ` ce denotes the standard pixel-wise cross-entropy loss . We define the overall loss function as : L = Ls + ↵Lu , ( 3 ) where ↵ controls the contribution of the unsupervised loss . 3.3 Adaptive Equalization Learning . The baseline framework , though achieving competitive results compared with previous related works , neglects the key issues in semi-supervised semantic segmentation . Due to the limited labeled data , semi-supervised learning tends to have poor performance on some certain categories , e.g. , tailed categories in Cityscapes dataset which exhibits a long-tailed label distribution . Insufficient training on these categories introduces more noise of pseudo labels which can disrupt the learning process . The proposed AEL framework aims to alleviate the degradation of under-performing categories during the semi-supervised training . Concretely , we maintain a confidence bank to record the performance of each category during training . The confidence bank enables us to identify the under-performing categories . To improve the performance of these categories and further make the training unbiased , we propose a series of technologies to efficiently leverage both labeled and unlabeled data , namely adaptive CutMix , adaptive Copy-Paste , adaptive equalization sampling and dynamic re-weighting . Confidence Bank . To tackle the biased training , previous methods [ 25 , 26 , 22 , 28 ] always rely on the prior knowledge such as the number of training samples of each category to design the ad hoc sampling and weighting strategies . However , the performance of each category is not always strictly proportional to the number of training samples , because some categories tend to have discriminative features and thus fewer samples are required for training . Inspired by the recent progress [ 29 ] which applies active learning on semantic segmentation , we propose to maintain a confidence bank to record the category-wise performance during training . An indicator is needed to assess the performance of each category . We consider several indicators , namely Confidence , Margin and Entropy . Formally , we define Confidence indicator as : Confc = 1 Nl NlX i=1 1 N c i N c iX j=1 p c ij , c 2 { 1 , . . . , C } ( 4 ) where C is the category number , N c i denotes the number of pixels belonging to category c according to its ground-truth yi , pcij denotes the c-th channel prediction of the j-th pixel in the i-th image . Define Margin indicator as : Marginc = 1 Nl NlX i=1 1 N c i N c iX j=1 ( pc ij max2 c02 { 1 , ... , C } p c 0 ij ) , c 2 { 1 , . . . , C } ( 5 ) where max2 ( · ) denotes the second largest value operator . At last , we define Entropy indicator as : Entc = 1 Nl NlX i=1 1 N c i N c iX j=1 CX c0=1 p c 0 ij log pc 0 ij , c 2 { 1 , . . . , C } . ( 6 ) For all of the indicators , we only take into account predictions from labeled data . Experimentally , the confidence indicator serves best in our AEL and thus we adopt it by default ( see Section 4.3 for the comparison ) . We use EMA to update the category-wise confidence at each training step : Confc k ⌧Confc k 1 + ( 1 ⌧ ) Confck , c 2 { 1 , . . . , C } , ( 7 ) where k denotes the k-th iteration , ⌧ 2 [ 0 , 1 ) is the momentum coefficient which is set to 0.999 experimentally . Through the confidence bank , we can easily identify the under-performing categories for the current model . Adaptive CutMix . Here we introduce the proposed adaptive CutMix ( see Figure 2 for illustration ) which is applied on the unlabeled data . It aims to increase the frequency of occurrence of the under-performing samples from the unlabeled data . We first formulate the original CutMix [ 27 ] as : Î = CutMix ( Crop ( I1 ) , I2 ) , ( 8 ) where I1 and I2 denote randomly selected unlabeled images , Î is the augmented image , and Crop ( · ) represents the random crop operation . Different from the original CutMix where unlabeled images are randmoly selected , the proposed adaptive CutMix gives under-performing categories a higher sampling probability . Specifically , we first convert the category-wise confidence stored in the confidence bank to the normalized sampling probability r 2 RC , which can be formulated as : r = Softmax ( 1 Conf ) . ( 9 ) According to the sampling probability , we randomly select an unlabeled image containing the sampled category as I1 , and another unlabeled image from the training batch is randomly selected as I2 . The Crop ( · ) operation is performed on the region containing the chosen category . After that , we can generate the augmented image by Eq 8 . Since the adaptive CutMix is performed on the unlabeled data without any annotations , we use predictions as approximate ground-truth , which works well in practice . Adaptive Copy-Paste . Copy-Paste [ 30 ] is an effective data augmentation strategy for instance segmentation . It yields significant gains on the challenging LVIS benchmark [ 31 ] , especially for rare object categories . The key idea behind the Copy-Paste augmentation is to paste objects from the source image to the target image . Inspired by this , we further propose the adaptive Copy-Paste ( see Figure 2 for illustration ) for semi-supervised semantic segmentation . Different from adaptive CutMix , adaptive Copy-Paste augmentation strives for efficiently leveraging the labeled data . Similarly , we involve confidence bank to assess category-wise performance and use Eq 9 to compute sampling probability . The under-performing categories have higher probability to be selected for Copy-Paste . Experimentally , the proposed adaptive Copy-Paste augmentation yields slightly better performance in the category level than the instance level . Thus we copy all pixels belonging to the sampled category in the source image and paste them on the target image . Following [ 30 ] , the augmented image is composed of two randomly selected images from the labeled data and a large scale jittering is applied . Adaptive Equalization Sampling . As described in Section 1 , due to the limited and unbalanced labeled data , the training tends to be biased . To alleviate the training bias , we propose a novel adaptive equalization sampling strategy which focuses training on a sparse set of under-performing samples and prevents the vast number of well-trained samples from overwhelming the model during training . Concretely , we define the sampling rate sc for category c as : s c = 1 Confc maxc2 { 1 , ... , C } ( 1 Confc ) , c 2 { 1 , . . . , C } , ( 10 ) where denotes a tunable parameter . Instead of using all pixels to compute the unsupervised loss , for category c with the sampling rate sc , we randomly sample a subset of pixels according to their predictions . Then the unsupervised loss in Eq 2 can be reformulated as : Lu = 1 Nu NuX i=1 1 P WH j=1 ij WHX j=1 ` ce ( ŷij , pij ) ij , ( 11 ) where ij = 1 indicates that the j-th pixel in the i-th image is sampled according to the sampling rate , otherwise ij is set to 0 , the other terms are the same as in Eq 2 . Dynamic Re-Weighting . The performance of the model depends on the quality of pseudo labels . Existing methods [ 2 , 4 , 19 ] usually adopt a higher threshold on classification score to filter out most of the pixels with low-confidence . Though this strategy could alleviate the noise raised by pseudolabeling , the strict criteria leads to lower recall for the pixels from under-performing categories , which hinders the training . Another option is to discard the threshold and involve all pixels into the training . However , much more noise is introduced simultaneously . To alleviate this issue , we propose a dynamic re-weighting strategy which adds a modulating factor to the unsupervised loss in the way of semi-supervised learning . On the basis of Eq 11 , we formulate our final unsupervised loss as : Lu = 1 Nu NuX i=1 1 P WH j=1 wij WHX j=1 wij ` ce ( ŷij , pij ) , ( 12 ) wij = max c2 { 1 , ... , C } ( pc ij ) ij , ( 13 ) where is the tunable parameter . Different from the Focal Loss [ 32 ] where the modulating factor is used for reducing the loss contribution from easy samples , our formulation aims to allocate more contributions for the convincing samples . The combination of adaptive equalization sampling and dynamic re-weighting not only involves more samples from the under-performing categories into the training , but also alleviate the noise raised by pseudo-labeling .

Summary: This paper tackles the challenging semi-supervised semantic segmentation task, with a focus on dealing with the imbalanced/long-tailed distribution. The authors propose two adaptive data augmentation strategies, one adaptive sampling schedule, and one re-weighting method to handle the long-tailed issue. The proposed method works well on two public standard benchmark datasets.