On Compressing U-net Using Knowledge Distillation

We study the use of knowledge distillation to compress the U-net architecture. We show that, while standard distillation is not sufficient to reliably train a compressed U-net, introducing other regularization methods, such as batch normalization and class re-weighting, in knowledge distillation significantly improves the training process. This allows us to compress a U-net by over 1000x, i.e., to 0.1% of its original number of parameters, at a negligible decrease in performance.

[1]  Mathieu Salzmann,et al.  Compression-aware Training of Deep Networks , 2017, NIPS.

[2]  Rich Caruana,et al.  Model compression , 2006, KDD '06.

[3]  Song Han,et al.  Trained Ternary Quantization , 2016, ICLR.

[4]  Andrew Zisserman,et al.  Very Deep Convolutional Networks for Large-Scale Image Recognition , 2014, ICLR.

[5]  Yann LeCun,et al.  Optimal Brain Damage , 1989, NIPS.

[6]  Geoffrey E. Hinton,et al.  Distilling the Knowledge in a Neural Network , 2015, ArXiv.

[7]  Pascal Fua,et al.  Supervoxel-Based Segmentation of Mitochondria in EM Image Stacks With Learned Shape Features , 2012, IEEE Transactions on Medical Imaging.

[8]  Chong Wang,et al.  Model Distillation with Knowledge Transfer from Face Classification to Alignment and Verification , 2017, 1709.02929.

[9]  Ali Farhadi,et al.  XNOR-Net: ImageNet Classification Using Binary Convolutional Neural Networks , 2016, ECCV.

[10]  Thomas Brox,et al.  U-Net: Convolutional Networks for Biomedical Image Segmentation , 2015, MICCAI.

[11]  Tony X. Han,et al.  Learning Efficient Object Detection Models with Knowledge Distillation , 2017, NIPS.

[12]  Xiaogang Wang,et al.  Face Model Compression by Distilling Knowledge from Neurons , 2016, AAAI.

[13]  Sergey Ioffe,et al.  Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift , 2015, ICML.

[14]  Jia Deng,et al.  Stacked Hourglass Networks for Human Pose Estimation , 2016, ECCV.

[15]  Song Han,et al.  Learning both Weights and Connections for Efficient Neural Network , 2015, NIPS.