Semantic Segmentation on PyTorch (include FCN, PSPNet, Deeplabv3, Deeplabv3+, DANet, DenseASPP, BiSeNet, EncNet, DUNet, ICNet, ENet, OCNet, CCNet, PSANet, CGNet, ESPNet, LEDNet, DFANet)

Semantic Segmentation on PyTorch

This project aims at providing a concise, easy-to-use, modifiable reference implementation for semantic segmentation models using PyTorch.

Installation

# semantic-segmentation-pytorch dependencies
pip install ninja tqdm

# follow PyTorch installation in https://pytorch.org/get-started/locally/
conda install pytorch torchvision -c pytorch

# install PyTorch Segmentation
git clone https://github.com/Tramac/awesome-semantic-segmentation-pytorch.git

# the following will install the lib with symbolic links, so that you can modify
# the files if you want and won't need to re-build it
cd awesome-semantic-segmentation-pytorch/core/nn
python setup.py build develop

Usage

Train

• Single GPU training

# for example, train fcn32_vgg16_pascal_voc:
python train.py --model fcn32s --backbone vgg16 --dataset pascal_voc --lr 0.0001 --epochs 50

• Multi-GPU training

# for example, train fcn32_vgg16_pascal_voc with 4 GPUs:
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS train.py --model fcn32s --backbone vgg16 --dataset pascal_voc --lr 0.0001 --epochs 50

Evaluation

• Single GPU evaluating

# for example, evaluate fcn32_vgg16_pascal_voc
python eval.py --model fcn32s --backbone vgg16 --dataset pascal_voc

• Multi-GPU evaluating

# for example, evaluate fcn32_vgg16_pascal_voc with 4 GPUs:
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS eval.py --model fcn32s --backbone vgg16 --dataset pascal_voc

Demo

cd ./scripts
python demo.py --model fcn32s_vgg16_voc --input-pic ./datasets/test.jpg

.{SEG_ROOT}
├── scripts
│   ├── demo.py
│   ├── eval.py
│   └── train.py

Support

Model

• FCN
• ENet
• PSPNet
• ICNet
• DeepLabv3
• DeepLabv3+
• DenseASPP
• EncNet
• BiSeNet
• PSANet
• DANet
• OCNet
• CGNet
• ESPNetv2
• CCNet
• DUNet(DUpsampling)
• FastFCN(JPU)
• LEDNet
• Fast-SCNN
• LightSeg
• DFANet

DETAILS for model & backbone.

.{SEG_ROOT}
├── core
│   ├── models
│   │   ├── bisenet.py
│   │   ├── danet.py
│   │   ├── deeplabv3.py
│   │   ├── deeplabv3+.py
│   │   ├── denseaspp.py
│   │   ├── dunet.py
│   │   ├── encnet.py
│   │   ├── fcn.py
│   │   ├── pspnet.py
│   │   ├── icnet.py
│   │   ├── enet.py
│   │   ├── ocnet.py
│   │   ├── ccnet.py
│   │   ├── psanet.py
│   │   ├── cgnet.py
│   │   ├── espnet.py
│   │   ├── lednet.py
│   │   ├── dfanet.py
│   │   ├── ......

Dataset

You can run script to download dataset, such as:

cd ./core/data/downloader
python ade20k.py --download-dir ../datasets/ade

.{SEG_ROOT}
├── core
│   ├── data
│   │   ├── dataloader
│   │   │   ├── ade.py
│   │   │   ├── cityscapes.py
│   │   │   ├── mscoco.py
│   │   │   ├── pascal_aug.py
│   │   │   ├── pascal_voc.py
│   │   │   ├── sbu_shadow.py
│   │   └── downloader
│   │       ├── ade20k.py
│   │       ├── cityscapes.py
│   │       ├── mscoco.py
│   │       ├── pascal_voc.py
│   │       └── sbu_shadow.py

Result

• PASCAL VOC 2012

Note: lr=1e-4, batch_size=4, epochs=80.

Overfitting Test

See TEST for details.

.{SEG_ROOT}
├── tests
│   └── test_model.py

To Do

• add train script
• remove syncbn
• train & evaluate
• test distributed training
• fix syncbn (Why SyncBN?)
• add distributed (How DIST?)

References

• PyTorch-Encoding
• maskrcnn-benchmark
• gloun-cv
• imagenet