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 work addresses the task of semi-supervised learning (SSL) in semantic segmentation. Following recent SOTAs in SSL, this work also advocates for the use of pseudo-labels on unlabeled data and heavy data augmentation. The main novelty of this work is the novel way to construct higher-quality pseudo-labels: besides the pixel-wise classifier's probabilistic outputs, the authors leverage as well CAM-based activation maps, named as SGC, as an additional pseudo-label source.  The final set of pseudo-labels is determined by linear combining the two soft pseudo-label sources with temperature adjustment. The authors conducted extensive experiments with lots of ablation studies to validate the proposed framework.