[论文解读] Training convolutional neural networks with megapixel images

Hans Pinckaers, Geert Litjens|arXiv (Cornell University)|Apr 16, 2018

AI in cancer detection参考文献 2被引用 8

一句话总结

本文提出流式随机梯度下降（SSGD），一种通过在图像块上进行部分前向传播以重建中间激活图，从而在显著减少GPU显存使用的情况下训练卷积神经网络的方法。该方法使使用6400万像素图像进行训练时，相比传统SGD显存使用量减少97%，在保持模型精度的同时实现数值等价的结果。

ABSTRACT

To train deep convolutional neural networks, the input data and the intermediate activations need to be kept in memory to calculate the gradient descent step. Given the limited memory available in the current generation accelerator cards, this limits the maximum dimensions of the input data. We demonstrate a method to train convolutional neural networks holding only parts of the image in memory while giving equivalent results. We quantitatively compare this new way of training convolutional neural networks with conventional training. In addition, as a proof of concept, we train a convolutional neural network with 64 megapixel images, which requires 97% less memory than the conventional approach.

研究动机与目标

• 解决在训练高分辨率图像（如十亿像素的组织病理学切片或卫星图像）时的显存瓶颈问题。
• 实现在不需大量GPU显存的情况下对全分辨率百万像素图像进行训练。
• 开发一种显存高效的方法，同时保持与标准SGD的数值等价性。
• 证明使用仅7 GB GPU显存即可训练6400万像素图像，相比标准SGD的235 GB，显存使用量大幅降低。
• 通过减少显存占用，使在标准硬件上训练大型网络或高分辨率数据成为可能。

提出的方法

• 通过在图像块上执行部分前向传播来重建中间激活图，随后将重建结果重新组合以继续网络的其余部分。
• 利用前向传播（P_forw）和反向传播（P_back）的公式计算块之间的重叠区域，以确保重建的准确性。
• 公式1用于计算前向传播的块重叠：P_forw = (k - s) + b((z - k) mod s)，其中k为卷积核大小，s为步长，z为图像尺寸，b为布尔标志。
• 公式2用于计算反向传播的更大重叠：P_back = 2(k - s) + b((z - k) mod s)，以考虑重叠区域之间的梯度依赖关系。
• 反向传播按块进行，梯度通过部分前向传播重建，避免存储所有激活值。
• 小批量梯度在图像间求和并取平均，同时通过在块边界保留边缘像素来校正池化/卷积操作带来的边缘效应。

实验结果

研究问题

• RQ1 能否开发一种训练方法，使全百万像素图像输入在不超出GPU显存限制的情况下进行？
• RQ2 是否可能实现标准SGD与显存减少型训练方法（如SSGD）之间的数值等价？
• RQ3 在反向传播过程中通过分块前向传播重建激活图，其显存与计算开销之间的权衡如何？
• RQ4 能否使用SSGD成功训练6400万像素图像？其最终显存占用量是多少？
• RQ5 SSGD在多大程度上限制了网络架构选择，例如是否支持批归一化或全局操作？

主要发现

• 标准SGD与SSGD之间的损失差异仅为1×10⁻⁴，证实了数值等价性。
• 测试准确率在标准SGD和SSGD下均为77%，表明性能相当。
• 在6400万像素图像上训练的网络达到78%的测试准确率，证明了大图像训练的可行性。
• SSGD将GPU显存使用量从标准SGD（小批量大小为8）的235 GB降低至7 GB，显存减少97%。
• 该方法在反向传播过程中需重新计算，引入了显著的计算开销。
• 依赖完整激活图的网络架构（如批归一化）与SSGD不兼容，因其依赖完整的特征图。