PseudoSeg: Designing Pseudo Labels for Semantic Segmentation

1 INTRODUCTION . Image semantic segmentation is a core computer vision task that has been studied for decades . Compared with other vision tasks , such as image classification and object detection , human annotation of pixel-accurate segmentation is dramatically more expensive . Given sufficient pixellevel labeled training data ( i.e. , high-data regime ) , the current state-of-the-art segmentation models ( e.g. , DeepLabv3+ ( Chen et al. , 2018 ) ) produce satisfactory segmentation prediction for common practical usage . Recent exploration demonstrates improvement over high-data regime settings with large-scale data , including self-training ( Chen et al. , 2020a ; Zoph et al. , 2020 ) and backbone pretraining ( Zhang et al. , 2020a ) . In contrast to the high-data regime , the performance of segmentation models drop significantly , given very limited pixel-labeled data ( i.e. , low-data regime ) . Such ineffectiveness at the low-data regime hinders the applicability of segmentation models . Therefore , instead of improving high-data regime segmentation , our work focuses on data-efficient segmentation training that only relies on few pixellabeled data and leverages the availability of extra unlabeled or weakly annotated ( e.g. , image-level ) data to improve performance , with the aim of narrowing the gap to the supervised models trained with fully pixel-labeled data . Our work is inspired by the recent success in semi-supervised learning ( SSL ) for image classification , demonstrating promising performance given very limited labeled data and a sufficient amount of unlabeled data . Successful examples include MeanTeacher ( Tarvainen & Valpola , 2017 ) , UDA ( Xie et al. , 2019 ) , MixMatch ( Berthelot et al. , 2019b ) , FeatMatch ( Kuo et al. , 2020 ) , and FixMatch ( Sohn et al. , 2020a ) . One outstanding idea in this type of SSL is consistency training : making predictions consistent among multiple augmented images . FixMatch ( Sohn et al. , 2020a ) shows that using high-confidence one-hot pseudo labels obtained from weakly-augmented unlabeled data to train strongly-augmented counterpart is the key to the success of SSL in image classification . ∗Work done during internship at Google Cloud AI Research . However , effective pseudo labels and well-designed data augmentation are non-trivial to satisfy for semantic segmentation . Although we observe that many related works explore the second condition ( i.e. , augmentation ) for image segmentation to enable consistency training framework ( French et al. , 2020 ; Ouali et al. , 2020 ) , we show that a wise design of pseudo labels for segmentation has great veiled potentials . In this paper , we propose PseudoSeg , a one-stage training framework to improve image semantic segmentation by leveraging additional data either with image-level labels ( weakly-labeled data ) or without any labels . PseudoSeg presents a novel design of pseudo-labeling to infer effective structured pseudo labels of additional data . It then optimizes the prediction of strongly-augmented data to match its corresponding pseudo labels . In summary , we make the following contributions : • We propose a simple one-stage framework to improve semantic segmentation by using a limited amount of pixel-labeled data and sufficient unlabeled data or image-level labeled data . Our framework is simple to apply and therefore network architecture agnostic . • Directly applying consistency training approaches validated in image classification renders particular challenges in segmentation . We first demonstrate how well-calibrated soft pseudo labels obtained through wise fusion of predictions from diverse sources can greatly improve consistency training for segmentation . • We conduct extensive experimental studies on the PASCAL VOC 2012 and COCO datasets . Comprehensive analyses are conducted to validate the effectiveness of this method at not only the low-data regime but also the high-data regime . Our experiments study multiple important open questions about transferring SSL advances to segmentation tasks . 2 RELATED WORK . Semi-supervised classification . Semi-supervised learning ( SSL ) aims to improve model performance by incorporating a large amount of unlabeled data during training . Consistency regularization and entropy minimization are two common strategies for SSL . The intuition behind consistencybased approaches ( Laine & Aila , 2016 ; Sajjadi et al. , 2016 ; Miyato et al. , 2018 ; Tarvainen & Valpola , 2017 ) is that , the model output should remain unchanged when the input is perturbed . On the other hand , the entropy minimization strategy ( Grandvalet & Bengio , 2005 ) argues that the unlabeled data can be used to ensured classes are well-separated , which can be achieved by encouraging the model to output low-entropy predictions . Pseudo-labeling ( Lee , 2013 ) is one of the methods for implicit entropy minimization . Recently , holistic approaches ( Berthelot et al. , 2019b ; a ; Sohn et al. , 2020a ) combining both strategies have been proposed and achieved significant improvement . By redesigning the pseudo label , we propose an efficient one-stage semi-supervised learning framework of semantic segmentation for consistency training . Semi-supervised semantic segmentation . Collecting pixel-level annotations for semantic segmentation is costly and prone to error . Hence , leveraging unlabeled data in semantic segmentation is a natural fit . Early methods utilize a GAN-based model either to generate additional training data ( Souly et al. , 2017 ) or to learn a discriminator between the prediction and the ground truth mask ( Hung et al. , 2018 ; Mittal et al. , 2019 ) . Consistency regularization based approaches have also been proposed recently , by enforcing the predictions to be consistent , either from augmented input images ( French et al. , 2020 ; Kim et al. , 2020 ) , perturbed feature embeddings ( Ouali et al. , 2020 ) , or different networks ( Ke et al. , 2020 ) . Recently , Luo & Yang ( 2020 ) proposes a dual-branch training network to jointly learn from pixel-accurate and coarse labeled data , achieving good segmentation performance . To push the performance of state of the arts , iterative self-training approaches ( Chen et al. , 2020a ; Zoph et al. , 2020 ; Zhu et al. , 2020 ) have been proposed . These methods usually assume the available labeled data is enough to train a good teacher model , which will be used to generate pseudo labels for the student model . However , this condition might not satisfy in the low-data regime . Our proposed method , on the other hand , realizing the ideas of both consistency regularization and pseudo-labeling in segmentation , consistently improves the supervised baseline in both low-data and high-data regimes . Weakly-supervised semantic segmentation . Instead of supervising network training with accurate pixel-level labels , many prior works exploit weaker forms of annotations ( e.g. , bounding boxes ( Dai et al. , 2015 ) , scribbles ( Lin et al. , 2016 ) , image-level labels ) . Most recent approaches use imagelevel labels as the supervisory signal , which exploits the idea of class activation map ( CAM ) ( Zhou et al. , 2016 ) . Since the vanilla CAM only focus on the most discriminative region of objects , dif- ferent ways to refine CAM have been proposed , including partial image/feature erasing ( Hou et al. , 2018 ; Wei et al. , 2017 ; Li et al. , 2018 ) , using an additional saliency estimation model ( Oh et al. , 2017 ; Huang et al. , 2018 ; Wei et al. , 2018 ) , utilizing pixel similarity to propagate the initial score map ( Ahn & Kwak , 2018 ; Wang et al. , 2020 ) , or mining and co-segment the same category of objects across images ( Sun et al. , 2020 ; Zhang et al. , 2020b ) . While achieving promising results using the approaches mentioned above , most of them require a multi-stage training strategy . The refined score maps are optimized again using a dense-CRF model ( Krähenbühl & Koltun , 2011 ) , and then used as the target to train a separate segmentation network . On the other hand , we assume there exists a small number of fully-annotated data , which allows us to learn stronger segmentation models than general methods without needing pixel-labeled data . 3 THE PROPOSED METHOD . In analogous to SSL for classification , our training objective in PseudoSeg consists of a supervised loss Ls applied to pixel-level labeled data Dl , and a consistency constraint Lu applied to unlabeled data Du 1 . Specifically , the supervised loss Ls is the standard pixel-wise cross-entropy loss on the weakly augmented pixel-level labeled examples : Ls = 1 N × |Dl| ∑ x∈Dl N−1∑ i=0 CrossEntropy ( yi , fθ ( ω ( xi ) ) ) , ( 1 ) where θ represents the learnable parameters of the network function f and N denotes the number of valid labeled pixels in an image x ∈ RH×W×3 . yi ∈ RC is the ground truth label of a pixel i in H×W dimensions , and fθ ( ω ( xi ) ) ∈ RC is the predicted probability of pixel i , where C is the number of classes to predict and ω ( · ) denotes the weak ( common ) data augmentation operations used by Chen et al . ( 2018 ) . During training , the proposed PseudoSeg estimates a pseudo label ỹ ∈ RH×W×C for each stronglyaugmented unlabeled data x in Du , which is then used for computing the cross-entropy loss . The unsupervised objective can then be written as : Lu = 1 N × |Du| ∑ x∈Du N−1∑ i=0 CrossEntropy ( ỹi , fθ ( β ◦ ω ( xi ) ) ) , ( 2 ) where β ( · ) denotes a stronger data augmentation operation , which will be described in Section 3.2 . We illustrate the unlabeled data training branch in Figure 1 . 3.1 THE DESIGN OF STRUCTURED PSEUDO LABELS . The next important question is how to generate the desirable pseudo label ỹ . A straightforward solution is directly using the decoder output of a trained segmentation model after confidence threshold- 1For simplicity , here we illustrate the method with unlabeled data and then show it can be easily adapted to use image-level labeled data in Section 3.2. ing , as suggested by Sohn et al . ( 2020a ) ; Zoph et al . ( 2020 ) ; Xie et al . ( 2020 ) ; Sohn et al . ( 2020b ) . However , as we demonstrate later in the experiments , the generated pseudo hard/soft labels as well as other post-processing of outputs are barely satisfactory in the low-data regime , and thus yield inferior final results . To address this issue , our design of pseudo-labeling has two key insights . First , we seek for a distinct yet efficient decision mechanisms to compensate for the potential errors of decoder outputs . Second , wisely fusing multiple sources of predictions to generate an ensemble and better-calibrated version of pseudo labels . Starting with localization . Compared with precise segmentation , learning localization is a simpler task as it only needs to provide coarser-grained outputs than pixel level of objects in images . Based on this motivation , we improve decoder predictions from the localization perspective . Class activation map ( CAM ) ( Zhou et al. , 2016 ) is a popular approach to provide localization for class-specific regions . CAM-based methods ( Hou et al. , 2018 ; Wei et al. , 2017 ; Ahn & Kwak , 2018 ) have been successfully adopted to tackle a different weakly supervised semantic segmentation task from us , where they assume only image-level labels are available . In practice , we adopt a variant of class activation map , Grad-CAM ( Selvaraju et al. , 2017 ) in PseudoSeg . From localization to segmentation . CAM estimates the strength of classifier responses on local feature maps . Thus , an inherent limitation of CAM-based approaches is that it is prone to attending only to the most discriminative regions . Although many weakly-supervised segmentation approaches ( Ahn & Kwak , 2018 ; Ahn et al. , 2019 ; Sun et al. , 2020 ) aim at refining CAM localization maps to segmentation masks , most of them have complicated post-processing steps , such as dense CRF ( Krähenbühl & Koltun , 2011 ) , which increases the model complexity when used for consistency training . Here we present a computationally efficient yet effective refinement alternative , which is learnable using available pixel-labeled data . Although CAM only localizes partial regions of interests , if we know the pairwise similarities between regions , we can propagate the CAM scores from the discriminative regions to the rest unattended regions . Actually , it has been shown in many works that the learned high-level deep features are usually good at similarity measurements of visual objects . In this paper , we find hypercolumn ( Hariharan et al. , 2015 ) with a learnable similarity measure function works fairly effective . Given the vanilla Grad-CAM output for all C classes , which can be viewed as a spatially-flatten 2-D vector of weight m ∈ RL×C , where each row mi is the response weight per class for one region i . Using a kernel functionK ( · , · ) : RH×RH → R that measures element-wise similarity given feature h ∈ RH of two regions , the propagated score m̂i ∈ RC can be computed as follows m̂i = mi + L−1∑ j=0 eK ( Wkhi , Wvhj ) ∑L−1 k=0 e K ( Wkhi , Wvhk ) mj ·Wc . ( 3 ) The goal of this function is to train Θ = { Wk , Wv ∈ RH×H , Wc ∈ RC×C } in order to propagate the high value in m to all adjacent elements in the feature space RH ( i.e. , hypercolumn features ) to region i . Adding mi in equation 3 indicates the skip-connection . To compute propagated score for all regions , the operations in equation 3 can be efficiently implemented with self-attention dotproduct ( Vaswani et al. , 2017 ) . For brevity , we denote this efficient refinement process output as selfattention Grad-CAM ( SGC ) maps in RH×H×C . Figure 6 in Appendix A specifies the architecture . Calibrated prediction fusion . SGC maps are obtained from low-resolution feature maps . It is then resized to the desired output resolution , and thus not sufficient at delineating crisp boundaries . However , compared to the segmentation decoder , SGC is capable of generating more locally-consistent masks . Thus , we propose a novel calibrated fusion strategy to take advantage of both decoder and SCG predictions for better pseudo labels . Specifically , given a batch of decoder outputs ( pre-softmax logits ) p̂ = fθ ( ω ( x ) ) and SGC maps m̂ computed from weakly-augmented data ω ( x ) , we generate the pseudo labels ỹ by F ( p̂ , m̂ ) = Sharpen ( γ Softmax ( p̂ Norm ( p̂ , m̂ ) ) + ( 1− γ ) Softmax ( m̂ Norm ( p̂ , m̂ ) ) , T ) . ( 4 ) Two critical procedures are proposed to use here to make the fusion process successful . First , p̂ and m̂ are from different decision mechanisms and they could have very different degrees of overconfidence . Therefore , we introduce the operation Norm ( a , b ) = √∑|a| i ( a 2 i + b 2 i ) as a nor- malization factor . It alleviates the over-confident probability after softmax , which could unfavorably dominate the resulted γ-averaged probability . Second , the distribution sharpening operation Sharpen ( a , T ) i = a 1/T i / ∑C j a 1/T j adjusts the temperature scalar T of categorical distribution ( Berthelot et al. , 2019b ; Chen et al. , 2020b ) . Figure 2 illustrates the predictions from different sources . More importantly , we investigate the pseudo-labeling from a calibration perspective ( Section 4.3 ) , demonstrating that the proposed soft pseudo label ỹ leads to a better calibration metric comparing to other possible fusion alternatives , and justifying why it benefits the final segmentation performance . Training . Our final training objective contains two extra losses : a classification loss Lx , and a segmentation lossLsa . First , to compute Grad-CAM , we add a one-layer classification head after the segmentation backbone and a multi-label classification loss Lx . Second , as specified in Appendix A ( Figure 6 ) , SGC maps are scaled as pixel-wise probabilities using one-layer convolution followed by softmax in equation 3 . Learning Θ to predict SGC maps needs pixel-labeled data Dl . It is achieved by an extra segmentation loss Lsa between SGC maps of pixel-labeled data and corresponding ground truth . All the loss terms are jointly optimized ( i.e. , Lu + Ls + Lx + Lsa ) , while Lsa only optimizes Θ ( achieved by stopping gradient ) . See Figure 7 in the appendix for further details .

Summary: This paper focuses on the problem of semi-supervised semantic segmentation, where less pixel-level annotations are used to train the network. A new one-stage training framework is proposed to include the process of localization cue generation, pseudo label refinement and training of semantic segmentation. Inspire by recent success in the semi-supervised learning (SSL), a novel calibrated fusion strategy is proposed to incorporate the concept of consistency training with data augmentation into the framework. Experiments on PASCAL VOC and MSCOCO benchmarks validate the effectiveness of the proposed method.