RankMatch: Exploring the Better Consistency Regularization for Semi-supervised Semantic Segmentation

This is the official PyTorch implementation for our CVPR2024 paper "RankMatch: Exploring the Better Consistency Regularization for Semi-supervised Semantic Segmentation".

Abstract

The key lie in semi-supervised semantic segmentation is how to fully exploit substantial unlabeled data to improve the model’s generalization performance by resorting to constructing effective supervision signals. Most methods tend to directly apply contrastive learning to seek additional supervision to complement independent regular pixel-wise consistency regularization. However, these methods tend not to be preferred ascribed to their complicated designs, heavy memory footprints and susceptibility to confirmation bias. In this paper, we analyze the bottlenecks exist in contrastive learning-based methods and offer a fresh perspective on inter-pixel correlations to construct more safe and effective supervision signals, which is in line with the nature of semantic segmentation. To this end, we develop a coherent RankMatch network, including the construction of representative agents to model inter-pixel correlation beyond regular individual pixel-wise consistency, and further unlock the potential of agents by modeling inter-agent relationships in pursuit of rank-aware correlation consistency. Extensive experimental results on multiple benchmarks, including mitochondria segmentation, demonstrate that RankMatch performs favorably against state-of-the-art methods. Particularly in the low-data regimes, RankMatch achieves significant improvements.

Getting Started

Installation

cd RankMatch
conda create -n rankmatch python=3.10.4
conda activate rankmatch
pip install -r requirements.txt
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html

Pretrained Backbone

ResNet-50 | ResNet-101

├── ./pretrained
    ├── resnet50.pth
    └── resnet101.pth

Dataset

• Pascal: JPEGImages | SegmentationClass
• Cityscapes: leftImg8bit | gtFine

Please modify your dataset path in configuration files.

The groundtruth masks were preprocessed by UniMatch.

The final folder structure should look like this:

├── RankMatch
	├── pretrained
		└── resnet50.pth, ...
	└── rankmatch.py
	└── ...
├── data
	├── pascal
		├── JPEGImages
		└── SegmentationClass
	├── cityscapes
		├── leftImg8bit
		└── gtFine

Usage

# use torch.distributed.launch
sh scripts/train.sh <num_gpu> <port>

Citation

If you find this project useful, please consider citing:

@inproceedings{mai2024rankmatch,
  title={RankMatch: Exploring the Better Consistency Regularization for Semi-supervised Semantic Segmentation},
  author={Mai, Huayu and Sun, Rui and Zhang, Tianzhu and Wu, Feng},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={3391--3401},
  year={2024}
}

Acknowledgements

RankMatch is primarily based on UniMatch. We are grateful to their authors for open-sourcing their code.