Lightweight Model for Real-Time Semantic Segmentation

This project aims at providing the popular lightweight model implementations for real-time semantic segmentation.

Usage

---

Train

• Single GPU training

python train.py --model mobilenet --dataset citys --lr 0.01 --epochs 240

• Multi-GPU training

# for example, train mobilenet with 4 GPUs:
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS train.py --model mobilenet --dataset citys --lr 0.01 --epochs 240

Evaluation

• Single GPU evaluating

python eval.py --model mobilenet_small --dataset citys

• Multi-GPU evaluating

# for example, evaluate mobilenet with 4 GPUs:
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS eval.py --model mobilenet --dataset citys

Result

• Cityscapes

Where: crop_size=768, lr=0.01, epochs=80.

Backbone	OHEM	Params(M)	FLOPs(G)	CPU(fps)	GPU(fps)	mIoU/pixACC	Model
mobilenet	✘	5.31	4.48	0.81	77.11	0.463/0.901	GoogleDrive,BaiduCloud(ybsg)
mobilenet	✓	5.31	4.48	0.81	75.35	0.526/0.909	GoogleDrive,BaiduCloud(u2y2)
mobilenetv2	✓	4.88	4.04	0.49	49.40	0.613/0.930	GoogleDrive,BaiduCloud(q2g5)
mobilenetv3_small	✓	1.02	1.64	2.59	104.56	0.529/0.908	GoogleDrive,BaiduCloud(e7no)
mobilenetv3_large	✓	2.68	4.59	1.39	79.43	0.584/0.916	GoogleDrive,BaiduCloud(i60c)
shufflenet	✓	6.89	5.68	0.57	43.79	0.493/0.901	GoogleDrive,BaiduCloud(6fjh)
shufflenetv2	✓	5.24	4.33	0.72	57.71	0.528/0.914	GoogleDrive,BaiduCloud(7pi5)
igcv3	✓	4.86	4.04	0.34	29.70	0.573/0.923	GoogleDrive,BaiduCloud(qe4f)
efficientnet-b0	✓	6.63	2.60	0.33	30.15	0.492/0.903	GoogleDrive,BaiduCloud(phuy)

• Improve

Model	batch_size	epochs	crop_size	init_weight	optimizer	mIoU/pixACC
mobilenetv3_small	4	80	768	kaiming_uniform	SGD	0.529/0.908
mobilenetv3_small	4	160	768	kaiming_uniform	SGD	0.587/0.918
mobilenetv3_small	8	160	768	kaiming_uniform	SGD	0.553/0/913
mobilenetv3_small	4	80	1024	kaiming_uniform	SGD	0.557/0.914
mobilenetv3_small	4	80	768	xavier_uniform	SGD	0.550/0.911
mobilenetv3_small	4	80	768	kaiming_uniform	Adam	0.549/0.911
mobilenetv3_small	8	160	1024	xavier_uniform	SGD	0.612/0.920

Support

• MobileNet
• MobileNetV2
• MobileNetV3
• ShuffleNet
• ShuffleNetV2
• IGCV3
• EfficientNet

To Do

• improve performance
• optimize memory
• check efficientnet
• replace nn.SyncBatchNorm by nn.BatchNorm.convert_sync_batchnorm
• check find_unused_parameters in nn.parallel.DistributedDataParallel

References

• maskrcnn-benchmark
• mobilenetv3-segmentation
• awesome-semantic-segmentation-pytorch