EU-Net: Efficient Training of U-Net for Biomedical Image Segmentation
摘要
Deep learning has revolutionized many fields, including computer vision, natural language processing, medical imaging, etc., through its ability to process complex data and learn discriminative representations. However, training deep learning models often requires substantial computational resources, leading to high energy consumption and unignorable environmental impact. The urgency of developing training-efficient methods to alleviate such influence is increasing, especially as the scaling laws are widely adopted in various research domains, requiring increasing computation to train a stronger model. In this work, building on the principles of backpropagation sparsification, we propose EU-Net, an efficient way of training U-Net in terms of backpropagation FLOPs and time consumption, and demonstrate its efficacy on multiple biomedical image segmentation tasks. Unlike previous methods that rely on top-k selection to keep the most important gradients, our approach adopts Bernoulli random subsampling with an interleaved subsampling schedule, achieving a balance of gradient direction and computation efficiency. We also mathematically prove the guaranteed convergence of our method under common optimization assumptions. We validate our approach on 20 biomedical segmentation tasks, demonstrating its ability to handle diverse data distributions and scales while maintaining computational efficiency. Experimental results show that our approach reduces the backpropagation FLOPs by 40% and time consumption by 20% while keeping comparable and even improved model performance compared to normal training. Code is available at https://github.com/lujiazho/EU-Net .