Post-Train Adaptive U-Net for Image Segmentation

TL;DR

This paper introduces a novel Post-Train Adaptive U-Net that can be trained once and then dynamically adapted to various device capabilities at runtime, improving segmentation quality across different hardware without retraining.

Contribution

It is the first application of Post-Train Adaptive approach to image segmentation, enabling flexible adaptation of U-Net models after initial training.

Findings

Improved Dice score on CamVid dataset

Model can switch between 6 PTA configurations at runtime

All PTA configurations outperform original U-Net in quality

Abstract

Typical neural network architectures used for image segmentation cannot be changed without further training. 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 post-train configuration can be done at runtime on any inference device including mobile. Also, the PTA approach has allowed to improve image segmentation Dice score on the CamVid dataset. The final trained model can be switched at runtime between 6 PTA configurations, which differ by inference time and quality. Importantly, all of the configurations have better quality than the original U-Net (No PTA) model.