An efficient solution for semantic segmentation: ShuffleNet V2 with atrous separable convolutions

We present a computationally efficient approach to semantic segmentation, while achieving a high mean intersection over union (mIOU), 70.33% on Cityscapes challenge. The network proposed is capable of running real-time on mobile devices.

Pre-print paper: https://arxiv.org/abs/1902.07476

If you find the code useful for your research, please consider citing us:

@article{turkmen2019efficient,
  title={An efficient solution for semantic segmentation: ShuffleNet V2 with atrous separable convolutions},
  author={T{\"u}rkmen, Sercan and Heikkil{\"a}, Janne},
  journal={arXiv preprint arXiv:1902.07476},
  year={2019}
}

Getting ready

1. Add tensorflow/models/slim to your python path in order to run most of the scripts! To do so follow these steps:
   1. Clone or download the tensorflow/models/slim repository to a separate folder.
   2. Add the path to the repository by running the following code: export PYTHONPATH=path_to_the_cloned_folder/tensorflow_models/research/slim:${PYTHONPATH}
2. Prepare dataset. Example scripts and code is available under the dataset folder. The dataset should be in tfrecord format.

Model zoo

Checkpoint name	Eval OS	Eval scales	Left-right Flip	mIOU	File Size
shufflenetv2_basic_cityscapes_67_7	16	[1.0]	No	67.7% (val)	4.9MB
shufflenetv2_dpc_cityscapes_71_3	16	[1.0]	No	71.3% (val)	6.3MB

Training

To learn more about the available flags you can check common.py and the specific script that you are trying to run (e.g. train.py).

Example training configuration

python train.py \
    --model_variant=shufflenet_v2 \
    --tf_initial_checkpoint=./checkpoints/model.ckpt \
    --training_number_of_steps=120000 \
    --base_learning_rate=0.001 \
    --fine_tune_batch_norm=True \
    --initialize_last_layer=False \
    --output_stride=16 \
    --train_crop_size=769 \
    --train_crop_size=769 \
    --train_batch_size=16 \
    --dataset=cityscapes \
    --train_split=train \
    --dataset_dir=./dataset/cityscapes/tfrecord \
    --train_logdir=./logs \
    --loss_function=sce

Example evaluation configuration

python evaluate.py \
    --model_variant=shufflenet_v2 \
    --eval_crop_size=1025 \
    --eval_crop_size=2049 \
    --output_stride=16 \
    --eval_logdir=./logs/eval \
    --checkpoint_dir=./logs \
    --dataset=cityscapes \
    --dataset_dir=./dataset/cityscapes/tfrecord

Exporting to TFLITE model

export_tflite.py script contains several parameters at the top of the script.

Running on Android

You can find an example script to run the this model and Tensorflow Lite interpreter for segmentation on Android in this repository.