UniMatch V2
This codebase contains the official PyTorch implementation of UniMatch V2:
UniMatch V2: Pushing the Limit of Semi-Supervised Semantic Segmentation Lihe Yang, Zhen Zhao, Hengshuang Zhao Preprint, 2024
TL;DR: We upgrade our UniMatch V1 by switching the outdated ResNet encoders to the most capable DINOv2 encoders. We unify the image-level and feature-level augmentations into a single learnable stream to challenge the powerful model. Based on this, we further design a Complementary Dropout to craft better dual views.
Results
We provide the training log of each reported value. You can refer to them during reproducing. We also provide all the checkpoints of our core experiments.
Pascal VOC 2012
| Method | Encoder | 1/16 (92) | 1/8 (183) | 1/4 (366) | 1/2 (732) | Full (1464) |
|---|---|---|---|---|---|---|
| UniMatch V1 | ResNet-101 | 75.2 | 77.2 | 78.8 | 79.9 | 81.2 |
| AllSpark | MiT-B5 | 76.1 | 78.4 | 79.8 | 80.8 | 82.1 |
| SemiVL | CLIP-Base | 84.0 | 85.6 | 86.0 | 86.7 | 87.3 |
| UniMatch V2 | DINOv2-Base | 86.3 | 87.9 | 88.9 | 90.0 | 90.8 |
Cityscapes
| Method | Encoder | 1/16 (186) | 1/8 (372) | 1/4 (744) | 1/2 (1488) |
|---|---|---|---|---|---|
| UniMatch V1 | ResNet-101 | 76.6 | 77.9 | 79.2 | 79.5 |
| AllSpark | MiT-B5 | 78.3 | 79.2 | 80.6 | 81.4 |
| SemiVL | CLIP-Base | 77.9 | 79.4 | 80.3 | 80.6 |
| UniMatch V2 | DINOv2-Base | 83.6 | 84.3 | 84.5 | 85.1 |
ADE20K
| Method | Encoder | 1/64 (316) | 1/32 (631) | 1/16 (1263) | 1/8 (2526) |
|---|---|---|---|---|---|
| UniMatch V1 | ResNet-101 | 21.6 | 28.1 | 31.5 | 34.6 |
| SemiVL | CLIP-Base | 33.7 | 35.1 | 37.2 | 39.4 |
| UniMatch V2 | DINOv2-Base | 38.7 | 45.0 | 46.7 | 49.8 |
COCO
| Method | Encoder | 1/512 (232) | 1/256 (463) | 1/128 (925) | 1/64 (1849) | 1/32 (3697) |
|---|---|---|---|---|---|---|
| UniMatch V1 | ResNet-101 | 31.9 | 38.9 | 44.4 | 48.2 | 49.8 |
| AllSpark | MiT-B5 | 34.1 | 41.7 | 45.5 | 49.6 | --- |
| SemiVL | CLIP-Base | 50.1 | 52.8 | 53.6 | 55.4 | 56.5 |
| UniMatch V2 | DINOv2-Base | 47.9 | 55.8 | 58.7 | 60.4 | 63.3 |
Real-World Large-Scale SSS Setting
In addition to the above traditional SSS settings, we also explore a real-world large-scale setting, where substantial images (e.g., 10K) have already been annotated, and menatime much more unlabeled images (e.g., 100K) are available. It is challenging but highly meaningful.
| Labeled Data (# Img) | + Unlabeled Data (# Img) | Improvement |
|---|---|---|
| COCO (118K) | COCO Extra (123K) | 66.4 → 67.1 |
| ADE20K (20K) | COCO Labeled (118K) | 54.1 → 54.9 |
| ADE20K (20K) | COCO All (118K + 123K) | 54.1 → 55.7 |
| Cityscapes (3K) | Cityscapes Extra (20K) | 85.2 → 85.5 |
Getting Started
Pre-trained Encoders
DINOv2-Small | DINOv2-Base | DINOv2-Large
├── ./pretrained
├── dinov2_small.pth
├── dinov2_base.pth
└── dinov2_large.pthDatasets
- Pascal: JPEGImages | SegmentationClass
- Cityscapes: leftImg8bit | gtFine
- ADE20K: images | annotations
- COCO: train2017 | val2017 | masks
Please modify your dataset path in configuration files.
The ADE20K and COCO annotations have already been pre-processed by us. You can use them directly.
├── [Your Pascal Path]
├── JPEGImages
└── SegmentationClass
├── [Your Cityscapes Path]
├── leftImg8bit
└── gtFine
├── [Your ADE20K Path]
├── images
│ ├── training
│ └── validation
└── annotations
├── training
└── validation
├── [Your COCO Path]
├── train2017
├── val2017
└── masksTraining
UniMatch V2
# use torch.distributed.launch
sh scripts/train.sh <num_gpu> <port>
# to fully reproduce our results, the <num_gpu> should be set as 4 on all four datasets
# otherwise, you need to adjust the learning rate accordingly
# or use slurm
# sh scripts/slurm_train.sh <num_gpu> <port> <partition>To train on other datasets or splits, please modify
dataset and split in train.sh.
FixMatch
Modify the method from 'unimatch_v2' to 'fixmatch' in train.sh.
Supervised Baseline
Modify the method from 'unimatch_v2' to 'supervised' in train.sh.
Citation
If you find this project useful, please consider citing:
@article{unimatchv2,
title={UniMatch V2: Pushing the Limit of Semi-Supervised Semantic Segmentation},
author={Yang, Lihe and Zhao, Zhen and Zhao, Hengshuang},
journal={arXiv:2410.10777},
year={2024}
}