ShelfNet

Link to project page
Link to my homepage
This is repository for my paper ShelfNet for real-time semantic segmentation, and avhieves both faster inference speed and higher segmentation accuracy, compared with other real-time models such as Lightweight-RefineNet.
This implementation is based on torch-encoding. Main difference is the structure of the model.

Results
We tested ShelfNet with ResNet50 and ResNet101 as the backbone respectively: they achieved 59 FPS and 42 FPS respectively on a GTX 1080Ti GPU with a 512x512 input image.
On PASCAL VOC 2012 test set, it achieved 84.2% mIoU with ResNet101 backbone and 82.8% mIoU with ResNet50 backbone.
It achieved 75.8% mIoU with ResNet50 backbone on Cityscapes dataset.

Requirements

Please refer to torch-encoding for implementation on synchronized batch-normalization layer.
PyTorch 0.4.1
Python 3.6
requests
nose
scipy
tqdm
Other requirements by torch-encoding.

Environment setup and data preparation

run python setup.py install to install torch-encoding
make sure you have the same path for a datset in /scripts/prepare_xx.py and /encoding/datasets/xxx.py, default path is ~/.encoding/data, which is a hidden folder. You will need to type Ctrl + h to show is in Files
run cd scripts
run python prepared_xx.py to prepare datasets, including MS COCO, PASCAL VOC, PASCAL Aug, PASCAL Context
Download test dataset from official evaluation server for PASCAL, extract and merge with training data folder, e.g. ~/.encoding/data/VOCdevkit

Configurations (refer to /experiments/option.py)

--model: which model to use, default is shelfnet, other options include pspnet, encnet,fcn
--backbone: backbone of the model, resnet50 or resnet101
--dataset: which dataset to train on, coco for MS COCO, pascal_aug for augmented PASCAL,pascal_voc for PASCAL VOC,pcontext for pascal context.
--aux: if type --aux, the model will use auxilliray layer, which is a FCN head based on the final block of backbone.
--se_loss: a context module based on final block of backbone, the shape is 1xm where m is number of categories. It penalizes whether a category is present or not.
--resume: default is None. It specifies the checkpoint to load
--ft: fine tune flag. If set as True, the code will resume from checkpoint but forget optimizer information.
--checkname: folder name to store trained weights
Other parameters are trevial, please refer to /experiments/segmentation/option.py for more details

Training scripts

run cd /experiments/segmentation
pre-train ShelfNet50 on COCO,
python train.py --backbone resnet50 --dataset coco --aux --se-loss --checkname ShelfNet50_aux
fine-tune ShelfNet50 on PASCAL_aug, you may need to double check the path for resume.
python train.py --backbone resnet50 --dataset pascal_aug --aux --se-loss --checkname ShelfNet50_aux --resume ./runs/coco/shelfnet/ShelfNet50_aux_se/model_best.pth.tar -ft
fine-tune ShelfNet50 on PASCAL VOC, you may need to double check the path for resume.
python train.py --backbone resnet50 --dataset pascal_voc --aux --se-loss --checkname ShelfNet50_aux --resume ./runs/pascal_aug/shelfnet/ShelfNet50_aux_se/model_best.pth.tar -ft

Test scripts

To test on PASCAL_VOC with multiple-scales input [0.5, 0.75, 1.0, 1.25, 1.5, 1.75].
python test.py --backbone resnet50 --dataset pascal_voc --resume ./runs/pascal_voc/shelfnet/ShelfNet50_aux_se/model_best.pth.tar
To test on PASCAL_VOC with single-scale input
python test_single_scale.py --backbone resnet50 --dataset pascal_voc --resume ./runs/pascal_voc/shelfnet/ShelfNet50_aux_se/model_best.pth.tar
Similar experiments can be performed on ShelfNet with ResNet101 backbone

Evaluation scripts

You can use the following script to generate ground truth - prediction pairs on PASCAL VOC validation set.
python evaluate_and_save.py --backbone resnet50 --dataset pascal_voc --resume ./runs/pascal_voc/shelfnet/ShelfNet50_aux_se/model_best.pth.tar --eval

Measure running speed

Measure running speed of ShelfNet on 512x512 image.
python test_speed.py --model shelfnet --backbone resnet101
python test_speed.py --model pspnet --backbone resnet101

Pre-trained weights