[论文解读] Deeper Image Quality Transfer: Training Low-Memory Neural Networks for 3D Images

Stefano B. Blumberg, Ryutaro Tanno+2 more|arXiv (Cornell University)|Aug 16, 2018

Advanced Image Processing Techniques参考文献 11被引用 7

一句话总结

该论文通过结合检查点技术和可逆网络（RevNets），提出了一种针对医学影像中深度3D卷积神经网络的内存高效训练框架，实现了仅4.0%的内存增长，同时显著提升网络深度。该方法在3D扩散磁共振成像（dMRI）超分辨率任务中将均方根误差（RMSE）降低了13%，超越了先前的最先进模型，同时在标准GPU上保持了低内存占用。

ABSTRACT

In this paper we address the memory demands that come with the processing of 3-dimensional, high-resolution, multi-channeled medical images in deep learning. We exploit memory-efficient backpropagation techniques, to reduce the memory complexity of network training from being linear in the network's depth, to being roughly constant $ - $ permitting us to elongate deep architectures with negligible memory increase. We evaluate our methodology in the paradigm of Image Quality Transfer, whilst noting its potential application to various tasks that use deep learning. We study the impact of depth on accuracy and show that deeper models have more predictive power, which may exploit larger training sets. We obtain substantially better results than the previous state-of-the-art model with a slight memory increase, reducing the root-mean-squared-error by $ 13\% $. Our code is publicly available.

研究动机与目标

• 解决在高分辨率医学影像上训练深度3D卷积网络时的高内存需求问题，该问题限制了标准GPU上模型的深度和性能。
• 克服GPU内存限制带来的妥协，例如使用小输入块或减小批量大小，这些妥协在医学图像深度学习中普遍存在。
• 通过将内存复杂度从线性降低至接近常数，实现对网络深度的内存高效训练。
• 在图像质量迁移（IQT）范式中展示该方法的有效性，用于增强低分辨率的扩散磁共振成像扫描。
• 表明通过内存效率实现的模型加深，可提升预测能力与泛化性能，尤其在使用更大训练集时优势更明显。

提出的方法

• 将可逆残差网络（RevNets）集成到3D卷积架构中，使激活值可从输出重建，从而将内存使用量降低至与深度无关的近似常数水平。
• 应用检查点技术，选择性地存储中间激活值，以计算成本的增加换取反向传播期间内存消耗的减少。
• 结合RevNets与检查点技术，手动控制前向与反向传播过程，实现深度网络的内存高效训练。
• 使用计算图抽象建模前向与反向传播，其中激活值按需重新计算，而非存储。
• 在基于ESPNet的3D图像超分辨率网络中实现该方法，通过在层间插入RevNet模块以扩展网络深度。
• 在单个NVIDIA Pascal GPU上使用标准优化策略进行模型训练，基于验证性能选择最佳模型，并在测试集上进行评估。

实验结果

研究问题

• RQ1 内存高效的反向传播技术是否能够在不显著增加内存开销的前提下，实现医学影像中更深的3D卷积网络训练？
• RQ2 通过内存减少实现的网络加深，是否能带来图像质量迁移性能的可测量提升？
• RQ3 检查点与RevNets的结合在多大程度上可降低3D医学影像网络的内存使用量，同时保持训练效率？
• RQ4 模型深度如何影响3D扩散磁共振成像超分辨率的性能？其性能是否从更大的训练集获益？
• RQ5 所提出的方法是否可推广至其他3D医学影像任务，如分割、图像合成与配准？

主要发现

• 与朴素反向传播相比，该方法将内存使用量降低了50%以上，且仅比基线ESPNet模型高出4.0%。
• 表现最佳的模型ESPNet-RN4在32名受试者上的RMSE相比先前最先进模型降低了13%（8.46 ± 0.67 vs. 9.76 ± 0.51）。
• 通过Wilcoxon符号秩检验确认该改进具有统计显著性（p = 0.0117，W = 0）。
• 更深的模型（如ESPNet-RN4）在脑部内部与外部区域均表现出更优性能，尤其在保留高频细节与组织边界方面。
• 该方法实现了内存增长可忽略的深层网络训练，使更大规模训练集得以更好利用，从而提升泛化能力。
• 尽管计算成本有所增加，所有模型均在单张GPU上24小时内完成训练，证明了其实际可行性。