结合多层特征及空间信息蒸馏的医学影像分割

Medical image segmentation based on multi-layer features and spatial information distillation

U-Net在医学影像分割领域是目前应用最广泛的分割模型,其“编码-解码”结构也成为了构建医学影像分割模型最常用的结构。尽管U-Net在许多领域实现了非常高的分割准确度,但是存在着计算复杂度高、推理速度慢、运行消耗内存大等问题,导致其难以在移动应用平台部署。为解决这一问题,提出了一种结合多层特征及空间信息蒸馏的医学影像分割方法TinyUnet。该方法使用轻量化的U-Net作为学生网络。考虑到小模型没有足够的学习能力,通过选择合适的蒸馏位置,对多层教师特征图进行蒸馏;同时加强教师网络深层特征图的边缘,并构建边缘关键点图结构,采用图卷积网络对学生网络进行空间信息蒸馏,从而补充重要的边缘信息和空间信息。实验表明:在3个医学影像数据集上,TinyUnet能够达到U-Net 98.3%~99.7%的分割准确度,但是将U-Net的参数量平均降低了99.6%,运算速度提高了约110倍;同时,与其他轻量化医学影像分割模型相比,TinyUnet不仅具有较高的分割准确度,而且占用内存更少,运行速度更快。

U-Net is currently the most widely used segmentation model,and its“coding-decoding”structure has also become the most commonly used structure for building medical image segmentation models.Although U-Net has achieved very high segmentation accuracy in many fields,but there are problems such as highcomputational complexity,slow reasoning speed,and high memory consumption,which makes it difficult to deploy on mobile application platforms.To solve this problem,a medical image segmentation method combining multi-layer features and spatial information distillation,named as TinyUnet,is proposedin this paper.This method uses the U-Net with fewer parameters as the student network,which is smaller and lighter than the original U-Net.Considering that the small model does not have enough learning ability,this method distils the multi-layer teacher feature maps by selecting the appropriate distillation position;at the same time,this method strengthens the edge of the deep feature map of the teacher network,constructs the edge key point map structure,and uses the graph convolution network to distil the spatial information of the student network,so as to guide the student network to obtain more effective edge information and spatial information.Experiments show that TinyUnet can maintain the segmentation accuracy of U-Net from 98.3%to 99.7%on the three medical datasets,but reduces the parameters of U-Net by 99.6%on average and increases the computing speed by about 110 times.Meanwhile,compared with other advanced compact medical image segmentation models,TinyUnet not only achieves good segmentation accuracy but also occupies less memory and runs faster.