semantic_segmentation

ADE20K Semantic segmentation with UniFormer

We currenent release the code and models for:

• Semantic FPN

• UperNet

Updates

05/22/2023

Lightweight models with Semantic FPN are released.

01/18/2022

1. Models with Semantic FPN are released..

2. Models with UperNet are released.

Model Zoo

The followed models and logs can be downloaded on Google Drive: total_models, total_logs.

We also release the models on Baidu Cloud: total_models (v5v2), total_logs (mhc6).

Note

• All the models are pretrained on ImageNet-1K without Token Labeling and Layer Scale. Reason can be found in issue #12.

Semantic FPN

We follow PVT to train our UniFormer with semantic FPN, please refer to fpn_seg.

Backbone	Lr Schd	mIoU	#params	FLOPs	Model	Log	Shell
UniFormer-XXS	80K	42.3	13.5M	-	google	google	run.sh/config
UniFormer-XS	80K	44.4	19.7M	-	google	google	run.sh/config
UniFormer-Sh14	80K	46.3	25M	172G	-	google	run.sh/config
UniFormer-Bh14	80K	47.0	54M	328G	google	google	run.sh/config

For the models with the window size of 32, we train one model but test with different test_config.

Backbone	Lr Schd	mIoU	#params	FLOPs	Model	Log	Shell
UniFormer-Sw32	80K	45.6	25M	183G	google	google	run.sh/config
UniFormer-Sh32	80K	46.2	25M	199G	google	google	run.sh/config
UniFormer-S	80K	46.6	25M	247G	google	google	run.sh/config
UniFormer-Bw32	80K	47.0	54M	310G	google	google	run.sh/config
UniFormer-Bh32	80K	47.7	54M	350G	google	google	run.sh/config
UniFormer-B	80K	48.0	54M	471G	google	google	run.sh/config

UperNet

We follow Swin Transformer to train our UniFormer with UperNet.

Backbone	Lr Schd	mIoU	MS mIoU	#params	FLOPs	Model	Log	Shell
UniFormer-Sh14	160K	46.9	48.0	52M	947G	-	google	run.sh/config
UniFormer-Bh14	160K	48.9	50.0	80M	1085G	-	google	run.sh/config

For the models with the window size of 32, we train one model but test with different test_config.

Backbone	Lr Schd	mIoU	MS mIoU	#params	FLOPs	Model	Log	Shell
UniFormer-Sw32	160K	46.6	48.4	52M	939G	google	google	run.sh/config
UniFormer-Sh32	160K	47.0	48.5	52M	955G	google	google	run.sh/config
UniFormer-S	160K	47.6	48.5	52M	1004G	google	google	run.sh/config
UniFormer-Bw32	160K	49.1	50.6	80M	1066G	google	google	run.sh/config
UniFormer-Bh32	160K	49.5	50.7	80M	1106G	google	google	run.sh/config
UniFormer-B	160K	50.0	50.8	80M	1227G	google	google	run.sh/config

Usage

Installation

Please refer to get_started for installation and dataset preparation.

Training

1. Download the pretrained models in our repository.

2. Simply run the training scripts in exp as followed:

   bash ./exp/upernet_global_small/run.sh

   Or you can train other models as follower:

   # single-gpu training
   python tools/train.py <CONFIG_FILE> --options model.backbone.pretrained_path=<PRETRAIN_MODEL> [other optional arguments]
   
   # multi-gpu training
   tools/dist_train.sh <CONFIG_FILE> <GPU_NUM> --options model.backbone.pretrained_path=<PRETRAIN_MODEL> [other optional arguments] 

[Note]:

• We use hybrid MHRA during testing and set the corresponding hyperparameters in the config.py:

  window: False, # whether use window MHRA
  hybrid: True, # whether use hybrid MHRA
  window_size: 32, # size of window (>=14)

• To avoid out of memory, we use torch.utils.checkpoint in the config.py:

  use_checkpoint=True, # whether use checkpoint
  checkpoint_num=[0, 0, 2, 0], # index for using checkpoint in every stage

Testing

# single-gpu testing
python tools/test.py <CONFIG_FILE> <DET_CHECKPOINT_FILE> --eval mIoU

# multi-gpu testing
tools/dist_test.sh <CONFIG_FILE> <DET_CHECKPOINT_FILE> <GPU_NUM> --eval mIoU

Acknowledgement

This repository is built based on mmsegmentation, Swin Transformer and PVT repository.