AutoSegEdge: Searching for the edge device real-time semantic segmentation based on multi-task learning

Real-time semantic segmentation is a challenging task for resource-constrained edge devices. We propose AutoSegEdge, based on Neural Architecture Search (NAS), a semantic segmentation approach that runs on edge devices in real-time. Besides accuracy, we employ FLOPs and latency on the target edge devices as search constraints. Our work is probably one of the first attempts to translate multi-objectives NAS into Multi-Task Learning. Be inspired by Multi-Task Learning, we regard the sub-objective in multi-objective NAS as a learning task in Multi-Task Learning. The total loss function of the multi-objective NAS is deconstructed into the weighted sum of the sub-objective loss function. However, the conflict among the sub-objective will cause the searched networks to “architecture collapse.” To avoid the multi-objectives NAS falls into “architecture collapse.” Based on uncertainty, this paper proposes a method to learn the weights of sub-objective loss functions automatically. AutoSegEdge was discovered from an efficient cell-level search space that integrates multi-resolution branches. Additionally, AutoSegEdge employs knowledge distillation to further boost accuracy. Finally, we accelerated AutoSegEdge using NVIDIA's TensorRT and deployed it on the Nvidia Jetson NX. Experiments demonstrate that multi-objectives NAS only requires 1.5 GPU days to obtain the best result on a single Nvidia Tesla V100 GPU. On the Cityscapes dataset, AutoSegEdge achieved an mIoU of 70.3% with 16.6 FPS on the Nvidia Jetson NX (and 194.54 FPS on an Nvidia Tesla V100 GPU) at the original resolution (1024 × 2048) using TensorRT. Our method is 2–3 × faster than existing state-of-the-art real-time methods while maintaining competitive accuracy. We also conducted robustness experiments to analyze our method and modules. The code is available: https://github.com/douziwenhit/AutoSeg_edge.git.