使用 Mesh-TensorFlow 进行超高分辨率图像分析

文 / 软件工程师 Le Hou 和 Youlong Cheng，Google Research

深度神经网络模型是构成大多数尖端图像分析和自然语言处理算法的基础。随着数据和模型并行计算等大规模深度学习技术的最新发展，我们已经可以使用包含数百万图像的数据集，在数分钟内即可训练完成大型 卷积神经网络 (CNN) 模型。

但在将 CNN 模型应用于  超分辨率图像 ( Ultra-High Resolution Image) 时依然充满挑战，例如 3D 计算机断层扫描 (CT) 图像（最多可包含 个像素）非常消耗资源。在现有技术条件下，处理器仍然需要托管至少 32GB 的局部、中间数据，而单个 GPU 或 TPU 的内存通常仅为 12-32GB。一种典型的解决方案是单独处理各个图像块，但由于信息丢失，实现过程复杂并会降低性能。

在与 Mayo Clinic 合作开展的“使用空间分区进行高分辨率医学图像分析(High Resolution Medical Image Analysis with Spatial Partitioning)” 项目中，我们利用 Mesh-TensorFlow 框架突破了大规模数据和模型并行计算的性能极限，并展示了该技术如何在不影响输入分辨率的前提下应用于超分辨率图像分析，从而确保实际可行性。

• Mayo Clinic
  https://www.mayoclinic.org/

• 使用空间分区进行高分辨率医学图像分析
  https://profs.etsmtl.ca/hlombaert/public/medneurips2019/28_CameraReadySubmission_MED-NeurIPS.pdf

• Mesh-TensorFlow
  https://papers.nips.cc/paper/8242-mesh-tensorflow-deep-learning-for-supercomputers.pdf

我们实现了 Halo Exchange 算法，该算法能够处理跨空间分区的卷积运算，以保留相邻分区之间的关系。基于此，我们便能使用 256 路模型并行计算在超分辨率图像（每个维度 512 像素的 3D 图像）上训练 3D U-Net。此外，我们还针对 GPU 和 TPU 开源了基于 Mesh-TensorFlow 的框架，供更多研究社区使用。

• 开源
  https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/experimental/model_executor.py
  

使用 Mesh-TensorFlow 实现数据和模型并行计算

我们的实现是基于 Mesh-TensorFlow 框架进行简单高效的数据和模型并行计算，使用户可以根据其定义的图像布局在设备网格中拆分张量。

• 张量
  https://developers.google.cn/machine-learning/glossary#tensor

例如，用户共计可以为 256 个处理器（每个处理器两个内核）提供 16 行 x 16 列的计算设备网格。然后，他们可以定义布局以将其图像的空间维度 x 映射到处理器行，空间维度 y 映射到处理器列，并将批次维度（即，要同时处理的图像片段数）映射到内核。训练批次的分区和分配由 Mesh-TensorFlow 在张量级别上实现，用户无需担心具体的实现细节。下图是对此概念的简化演示：

超分辨率图像的空间分区，本例中为 3D CT 扫描

使用 Halo Exchange 算法进行空间分区

在图像上执行的卷积运算通常会应用一个超出帧边缘的过滤器。尽管在处理单个图像时有多种方法可以解决此问题，但标准方法并未考虑到帧边缘之外的信息，对于分段图像而言可能仍然存在相关性。为了获得准确的结果，在对已进行空间分区、跨处理器重新分配的图像执行卷积运算时，我们必须考虑每个图像片段的相邻片段。

一种可能的解决方案是在每个空间分区中包含重叠区域。但是，由于出现多个后续卷积层的可能性极大，且每个卷积层都会加入重叠区，因此重叠区会比较大——实际上，在大多数情况下，重叠区可能会覆盖整个图像。此外，所有重叠区域均须从第一层就开始包含，但此做法可能会遇到内存限制，这也是我们正努力解决的问题。

考虑到这一点，我们采取了完全不同的解决方案，即执行一个名为 Halo Exchange 的数据通信步骤。在每次执行卷积运算之前，每个空间分区都会与其相邻分区交换（接收和发送）边距，从而有效拓宽图像片段边距。然后，即可在本地的每个设备上应用卷积运算。这可确保整个图像的卷积结果保持一致，无论是否进行空间分区都不受影响。

Halo Exchange 可确保跨分区的卷积正确处理图像片段边缘

概念验证—肝脏肿瘤 CT 扫描图像分割

随后，我们将此框架应用于肝脏肿瘤 3D CT 扫描图像分割任务（LiTS 测试平台）。针对评估指标，我们使用范围在 0.0 到 1.0 之间的 Sørensen–Dice 系数，得分为 0 表示分段区域与真实的肿瘤区域之间没有重叠，得分为 1 则表示两个区域完美匹配。下方结果表明，数据分辨率越高，获得的结果就越精准。尽管在使用全 分辨率（在 x、 y、 z 方向上均有 512 个像素）时返回值呈递减趋势，但这项研究确实为超分辨率图像分析提供了更多的可能性。

• LiTS
  https://arxiv.org/abs/1901.04056

数据分辨率越高，分割准确度就越高

结论

现有的数据和模型并行计算技术可以训练具有数十亿参数的神经网络，但无法处理像素数高于 ~ 的输入图像。在本次研究中，我们探索了 CNN 在这些超分辨率图像上的适用性，并向大家展示了可喜的结果。基于 Mesh-TensorFlow 对 GPU 和 TPU 均奏效实现的结果与发布的代码，我们希望为一些之前无法完成的任务提供可能的解决方案。

• 发布的代码
  https://github.com/tensorflow/mesh/tree/master/mesh_tensorflow/experimental

致谢

感谢 Mayo Clinic 的协作者 Panagiotis Korfiatis 博士和 Daniel Blezek 博士提供初始 3D U-net 模型和训练数据。感谢 Greg Mikels 与 Mayo Clinic 合作完成概念验证工作。最后特别感谢本篇论文的所有联合作者，尤其是 Noam Shazeer。

— 推荐阅读 —