Bootstrapping Semantic Segmentation with Regional Contrast

1 INTRODUCTION . Semantic segmentation is an essential part of applications such as scene understanding and autonomous driving , whose goal is to assign a semantic label to each pixel in an image . Significant progress has been achieved by use of large datasets with high quality human annotations . However , labelling images with pixel-level accuracy is time consuming and expensive ; for example , labelling a single image in CityScapes can take more than 90 minutes ( Cordts et al. , 2016 ) . When deploying semantic segmentation models in practical applications where only limited labelled data are available , high quality ground-truth annotation is a significant bottleneck . To reduce the need for labelled data , there is a recent surge of interest in leveraging unlabelled data for semi-supervised learning . Previous methods include improving segmentation models via adversarial learning ( Hung et al. , 2019 ; Mittal et al. , 2019 ) and self-training ( Zou et al. , 2019 ; 2018 ; Zhu et al. , 2020 ) . Others focus on designing advanced data augmentation strategies to generate pseudo image-annotation pairs from unlabelled images ( Olsson et al. , 2021 ; French et al. , 2020 ) . In both semi-supervised and supervised learning , a segmentation model often predicts smooth label maps , because neighbouring pixels are usually of the same class , and rarer high-frequency regions are typically only found in object boundaries . This learning bias produces blurry contours and regularly mis-labels rare objects . After carefully examining the label predictions , we further observe that wrongly labelled pixels are typically confused with very few other classes ; e.g . a pixel labelled as rider has a much higher chance of being wrongly classified as person , compared to train or bus . By understanding this class structure , learning can be actively focused on the challenging pixels to improve overall segmentation quality . Here we propose ReCo , a contrastive learning framework designed at a regional level . Specifically , ReCo is a new loss function which helps semantic segmentation not only to learn from local context ( neighbouring pixels ) , but also from global semantic class relationships across the entire dataset . ReCo performs contrastive learning on a pixel-level dense representation , as visualised in Fig . 1 . For each semantic class in a mini-batch , ReCo samples a set of pixel-level representations ( queries ) , and encourages them to be close to the class mean representation ( positive keys ) , and simultaneously pushes them away from representations sampled from other classes ( negative keys ) . For pixel-level contrastive learning with high-resolution images , it is impractical to sample all pixels . In ReCo , we actively sample a sparse set of queries and keys , consisting of less than 5 % of all available pixels . We sample negative keys from a learned distribution based on the relative distance between the mean representation of each negative key and the query class . This distribution can be interpreted as a pairwise semantic class relationship , dynamically updated during training . We sample queries for those having a low prediction confidence . Active sampling helps ReCo to rapidly focus on the most confusing pixels for each semantic class , and requires minimal additional memory . ReCo enables a high-accuracy segmentation model to be trained with very few human annotations . We evaluate ReCo in a semi-supervised setting , with two different modes : i ) Partial Dataset Full Labels — a sparse subset of training images , where each image has full ground-truth labels , and the remaining images are unlabelled ; ii ) Partial Labels Full Dataset — all images have some labels , but covering only a sparse subset of pixels within each image . In both settings , we show that ReCo can consistently improve performance across all methods and datasets . 2 RELATED WORK . Semantic Segmentation One recent direction is in designing more effective deep convolutional neural networks . Fully convolutional networks ( FCNs ) ( Long et al. , 2015 ) are the foundation of modern segmentation network design . They were later improved with dilated/atrous convolutions with larger receptive fields , capturing more long range information ( Chen et al. , 2017 ; 2018 ) . Alternative approaches include encoder-decoder architectures ( Ronneberger et al. , 2015 ; Kirillov et al. , 2019 ) , sometimes using skip connections ( Ronneberger et al. , 2015 ) to refine filtered details . A parallel direction is to improve optimisation strategies , by designing loss functions that better respect class imbalance ( Lin et al. , 2017 ) or using rendering strategy to refine uncertain pixels from high-frequency regions improving the label quality ( Kirillov et al. , 2020 ) . ReCo is built upon this line of research , to improve segmentation by providing additional supervision on hard pixels . Semi-supervised Classification and Segmentation The goal of semi-supervised learning is to improve model performance by taking advantage of a large amount of unlabelled data during training . Here consistency regularisation and entropy minimisation are two common strategies . The intuition is that the network ’ s output should be invariant to data perturbation and geometric transformation . Based on these strategies , many semi-supervised methods have been developed for image classification ( Sohn et al. , 2020 ; Tarvainen & Valpola , 2017 ; Berthelot et al. , 2019 ; Kuo et al. , 2020 ) . However , for segmentation , generating effective pseudo-labels and well-designed data augmentation are non-trivial . Some solutions improved the quality of pseudo-labelling , using adversarial learning ( Hung et al. , 2019 ; Mittal et al. , 2019 ) or enforcing consistency from different augmented images ( French et al. , 2020 ; Olsson et al. , 2021 ) . In this work , we show that we can improve the performance of current semi-supervised segmentation methods by jointly training with a suitable auxiliary task . Contrastive Learning Contrastive learning learns a similarity function to bring views of the same data closer in representation space , whilst pushing views of different data apart . Most recent contrastive frameworks learn similarity scores based on global representations of the views , parameterising data with a single vector ( He et al. , 2020 ; Chen et al. , 2020 ; Khosla et al. , 2020 ) . Dense representations , on the other hand , rely on pixel-level representations and naturally provide additional supervision , capturing fine-grained pixel correspondence . Contrastive pre-training based on dense representations has recently been explored , and shows better performance in dense prediction tasks , such as object detection and keypoint detection ( Wang et al. , 2021b ; O. Pinheiro et al. , 2020 ) . Contrastive Learning for Semantic Segmentation Contrastive learning has been recently studied to improve semantic segmentation , with a number of different design strategies . Zhang et al . ( 2021 ) and Zhao et al . ( 2021 ) both perform contrastive learning via pre-training , based on the generated auxiliary labels and ground-truth labels respectively , but at the cost of huge memory consump- tion . In contrast , ours performs contrastive learning whilst requiring much less memory , via active sampling . In concurrent work , ( Wang et al. , 2021a ; Alonso et al. , 2021 ) also perform contrastive learning with active sampling . However , whilst both these methods are applied to a stored feature bank , ours focuses on sampling features on-the-fly . Active sampling in Alonso et al . ( 2021 ) is further based on learnable , class-specific attention modules , whilst ours only samples features based on relation graphs and prediction confidence , without introducing any additional computation overhead , which results in a simpler and much more memory-efficient implementation . 3 RECO – REGIONAL CONTRAST . 3.1 PIXEL-LEVEL CONTRASTIVE LEARNING . Let ( X , Y ) be a training dataset with training images x ∈ X and their corresponding C-class pixellevel segmentation labels y ∈ Y , where y can be either provided in the original dataset , or generated automatically as pseudo-labels . A segmentation network f is then optimised to learn a mapping fθ : X 7→ Y , parameterised by network parameters θ . This segmentation network f can be decomposed into two parts : an encoder network : φ : X 7→ Z , and a decoder classification head ψc : Z 7→ Y . To perform pixel-level contrastive learning , we additionally attach a decoder representation head ψr on top of the encoder network φ , parallel to the classification head , mapping the encoded feature into a higher m-dimensional dense representation with the same spatial resolution as the input image : ψr : Z 7→ R , R ∈ Rm . This representation head is only applied during training to guide the classifier using the ReCo loss as an auxiliary task , and is removed during inference . A pixel-level contrastive loss is a function which encourages queries rq to be similar to the positive key r+k , and dissimilar to the negative keys r − k . All queries and keys are sampled from the decoder representation head : rq , r + , − k ∈ R. In ReCo , we use a pixel-level contrastive loss across all available semantic classes in each mini-batch , with the distance between keys and queries measured by their normalised dot product . The general formation of the ReCo loss Lreco is then defined as : Lreco = ∑ c∈C ∑ rq∼Rcq − log exp ( rq · rc , +k /τ ) exp ( rq · rc , +k /τ ) + ∑ r−k ∼R c k exp ( rq · r−k /τ ) , ( 1 ) for which C is a set containing all available classes in the current mini-batch , τ is the temperature control of the softness of the distribution , Rcq represents a query set containing all representations whose labels belong to class c , Rck represents a negative key set containing all representations whose labels do not belong to class c , and rc , +k represents the positive key which is the mean representation of class c. Suppose P is a set containing all pixel coordinates with the same resolution as R , these queries and keys are then defined as : Rcq = ⋃ [ u , v ] ∈P 1 ( y [ u , v ] = c ) r [ u , v ] , Rck = ⋃ [ u , v ] ∈P 1 ( y [ u , v ] 6= c ) r [ u , v ] , rc , +k = 1 |Rcq| ∑ rq∈Rcq rq . ( 2 ) 3.2 ACTIVE HARD SAMPLING ON QUERIES AND KEYS . Contrastive learning on all pixels in high-resolution images would be computationally expensive . Here , we introduce active hard sampling strategies to optimise only a sparse set of queries and keys . Active Key Sampling When classifying a pixel , a semantic network might be uncertain only over a very small number of candidates , among all available classes . The uncertainty from these candidates typically comes from a close spatial ( e.g . rider and bicycle ) or semantic ( e.g . horse and cow ) relationship . To reduce this uncertainty , we propose to sample negative keys non-uniformly , based on the relative distance between each negative key class and the query class . This involves building a pair-wise class relationship graph G , with G ∈ R|C|×|C| , computed and dynamically updated for each mini-batch . This pair-wise relationship is measured by the normalised dot product between the mean representation from a pair of two classes and is defined as : G [ p , q ] = ( rp , +k · r q , + k ) , ∀p , q ∈ C , and p 6= q . ( 3 ) We further apply SoftMax to normalise these pair-wise relationships among all negative classes j for each query class c , which produces a probabilistic distribution : exp ( G [ c , i ] ) / ∑ j∈C , j 6=c exp ( G [ c , j ] ) . We sample negative keys for each class i based on this distribution , to learn the corresponding query class c. This procedure allocates more samples to hard , confusing classes chosen specifically for each query class , helping the segmentation network to learn a more accurate decision boundary . Active Query Sampling Due to the natural class imbalance in semantic segmentation , it is easy to over-fit on common classes , such as the road and building classes in the CityScapes dataset , or the background class in the Pascal VOC dataset . These common classes contribute to the majority of pixel space in training images , and so randomly sampling queries will under-sample rare classes and provide minimal supervision to these classes . Therefore , we instead sample hard queries — for those whose corresponding pixel prediction confidence is below a defined threshold . Accordingly , ReCo ’ s loss would then guide the segmentation network by providing appropriate supervision on these less certain pixels . The easy and hard queries are defined as follows , and visualised in Fig . 2 , Rc , easyq = ⋃ rq∈Rcq 1 ( ŷq > δs ) rq , Rc , hardq = ⋃ rq∈Rcq 1 ( ŷq ≤ δs ) rq , ( 4 ) where ŷq is the predicted confidence of label c after the SoftMax operation corresponding to the same pixel location as rq , and δs is the user-defined confidence threshold .

Summary: This paper proposes ReCo, a regional contrastive learning method for semi-supervised semantic segmentation. The query and key pixel sampling methods are proposed for efficient learning. The proposed method showed state-of-the-art level performances in various settings and various datasets.