SegmenTron

PyTorch for Semantic Segmentation

Introduce

This repository contains some models for semantic segmentation and the pipeline of training and testing models, implemented in PyTorch.

Model zoo

Model	Backbone	Datasets	eval size	Mean IoU(paper)	Mean IoU(this repo)
DeepLabv3_plus	xception65	cityscape(val)	(1025,2049)	78.8	78.93
DeepLabv3_plus	xception65	coco(val)	480/520	-	70.50
DeepLabv3_plus	xception65	pascal_aug(val)	480/520	-	89.56
DeepLabv3_plus	xception65	pascal_voc(val)	480/520	-	88.39
DeepLabv3_plus	resnet101	cityscape(val)	(1025,2049)	-	78.27
Danet	resnet101	cityscape(val)	(1024,2048)	79.9	79.34
Pspnet	resnet101	cityscape(val)	(1025,2049)	78.63	77.00

real-time models

Model	Backbone	Datasets	eval size	Mean IoU(paper)	Mean IoU(this repo)	FPS
ICnet	resnet50(0.5)	cityscape(val)	(1024,2048)	67.8	-	41.39
DeepLabv3_plus	mobilenetV2	cityscape(val)	(1024,2048)	70.7	70.3	46.64
BiSeNet	resnet18	cityscape(val)	(1024,2048)	-	-	39.90
LEDNet	-	cityscape(val)	(1024,2048)	-	-	31.78
CGNet	-	cityscape(val)	(1024,2048)	-	-	46.11
HardNet	-	cityscape(val)	(1024,2048)	75.9	-	69.06
DFANet	xceptionA	cityscape(val)	(1024,2048)	70.3	-	21.46
HRNet	w18_small_v1	cityscape(val)	(1024,2048)	70.3	70.5	66.01
Fast_SCNN	-	cityscape(val)	(1024,2048)	68.3	68.9	145.77

FPS was tested on V100.

Environments

• python 3
• torch >= 1.1.0
• torchvision
• pyyaml
• Pillow
• numpy

INSTALL

python setup.py develop

if you do not want to run CCNet, you do not need to install, just comment following line in segmentron/models/__init__.py

from .ccnet import CCNet

Dataset prepare

Support cityscape, coco, voc, ade20k now.

Please refer to DATA_PREPARE.md for dataset preparation.

Pretrained backbone models

pretrained backbone models will be download automatically in pytorch default directory(~/.cache/torch/checkpoints/).

Code structure

├── configs    # yaml config file
├── segmentron # core code
├── tools      # train eval code
└── datasets   # put datasets here 

Train

Train with a single GPU

CUDA_VISIBLE_DEVICES=0 python -u tools/train.py --config-file configs/cityscapes_deeplabv3_plus.yaml

Train with multiple GPUs

CUDA_VISIBLE_DEVICES=0,1,2,3 ./tools/dist_train.sh ${CONFIG_FILE} ${GPU_NUM} [optional arguments]

Eval

Eval with a single GPU

You can download trained model from model zoo table above, or train by yourself.

CUDA_VISIBLE_DEVICES=0 python -u ./tools/eval.py --config-file configs/cityscapes_deeplabv3_plus.yaml \
TEST.TEST_MODEL_PATH your_test_model_path

Eval with a multiple GPUs

CUDA_VISIBLE_DEVICES=0,1,2,3 ./tools/dist_test.sh ${CONFIG_FILE} ${GPU_NUM} \
TEST.TEST_MODEL_PATH your_test_model_path

References

• PyTorch-Encoding
• detectron2
• gloun-cv