POST-TRAIN ADAPTIVE U-NET FOR IMAGE SEGMENTATION

Abstract

Many fields benefit from fast and accurate image segmentation. Convolutional neural networks show the best accuracy solving the task. Applications include medical or satellite imaging, autonomous driving, etc. Typical neural network architectures used for image segmentation are expected to be fully configured before the training procedure starts. To change the network architecture additional training steps are required. This is quite limiting as the network might not only be executed on a powerful server, but also on a mobile or edge device. Adaptive neural networks offer a solution to the problem by allowing certain adaptivity after the training process is complete. In this work for the first time, we apply Post-Train Adaptive (PTA) approach to the task of image segmentation. We introduce U-Net+PTA neural network, which can be trained once, and then adapted to different device performance categories. The two key components of the approach are PTA blocks and PTA-sampling training strategy. The PTA blocks were added into the U-Net neural network with MobileNetV2 backbone. The post-train configuration can be done at runtime on any inference device including, but not limited to mobile devices. In addition to post-train neural network configuration, the PTA approach has allowed to improve image segmentation quality (Dice score) on the CamVid dataset. The final trained model can be switched at runtime between 6 PTA configurations. These configurations differ by inference time and quality. Importantly, all of the configurations have better quality than the original U-Net (No PTA) model. The possible future research direction is to expand the inference time difference between heavy and light PTA configurations to allow a single trained PTA-based network to target even more device performance categories.