用于脑肿瘤分割的具有可逆移动倒置瓶颈的内存高效 3D U-Net
图像与视频处理
2021-04-22 v2 计算机视觉与模式识别
机器学习
摘要
我们提出将内存节省技术与传统 U-Net 架构相结合，以提升 Brain Tumor Segmentation (BraTS) 挑战赛中模型的复杂度。BraTS 挑战赛的任务是将 240x240x155x4 的输入图像进行 3D 分割为一组肿瘤类别。由于数据量大且需要 3D 卷积层，该任务非常耗费内存。为此，先前的方法使用较小的裁剪图像，同时限制模型的深度和宽度。我们的 3D U-Net 使用了在 MobileNetV2、MnasNet 以及较新的 EfficientNet 架构中定义的可逆版本移动倒置瓶颈块，以在训练期间节省激活内存。使用可逆层使模型能够在该层输出给定时重新计算输入激活，从而通过消除前向传播中存储激活的需求来节省内存。倒置残差瓶颈块使用轻量级深度可分离卷积，通过将卷积分解为逐点卷积和深度卷积来减少计算量。此外，该块通过将中间扩展层置于输入与输出线性 1x1 卷积之间，反转了传统瓶颈块，从而减少了总通道数。在固定的内存预算下，借助这些内存节省技术，我们能够训练比相应的非可逆网络大至多 3 倍的图像体积、深度增加 25% 的模型，或通道数至多 2 倍的模型。
引用
@article{arxiv.2104.09648,
title = {Memory Efficient 3D U-Net with Reversible Mobile Inverted Bottlenecks for Brain Tumor Segmentation},
author = {Mihir Pendse and Vithursan Thangarasa and Vitaliy Chiley and Ryan Holmdahl and Joel Hestness and Dennis DeCoste},
journal= {arXiv preprint arXiv:2104.09648},
year = {2021}
}
备注
11 pages, 5 figures, Published at MICCAI Brainles 2020