使用 Cloud TPU 处理大型图片

本文档介绍了如何使用 Cloud TPU 和 TPUEstimator 的空间分区功能对超大型图片和视频进行训练。

Cloud TPU 空间分区功能可以将单个图片拆分在多个 TPU 芯片上。利用空间分区功能,您可以轻松地扩缩模型,以便与过大而无法放入加速器芯片上的单核可使用的内存的输入图片一起运行。

此功能可在多个 TPU 核心之间分布激活,从而使您能够将模型扩缩为使用 2 个、4 个、8 个甚至是 16 个核心进行训练。如果扩缩可用于训练的处理器数量,则无需对图片数据进行降采样,从而提供更高的准确度和更高效的训练性能。

您可以将空间分区应用于以下任务:

  • 3D 计算机断层 (CT) 扫描图片分割
  • 视频内容分析
  • 在无需对您的图片数据进行降采样的情况下对自动驾驶进行对象检测。

启用 TPUEstimator 的空间分区功能

在 Tensorflow 中,Tensorflow 的 XLA 优化编译器会自动处理所有 Cloud TPU 核心之间的通信。无需更改代码即可为模型启用空间分区功能。由于 TPUEstimator 支持空间分区 API,因此您只需配置如何对 TPUConfig 中的每个输入张量进行分区即可。

图片

示例

以下代码展示了一个用于图片分类模型的四向空间分区的 TPUConfig 示例。张量会根据高度维度拆分为四个部分(假定张量具有[批量、高度、宽度、通道]的形状)。

tpu_config=tpu_config.TPUConfig(
    iterations_per_loop=100,
    num_cores_per_replica=4,
    per_host_input_for_training=tpu_config.InputPipelineConfig.PER_HOST_V2,
    input_partition_dims=[[1, 4, 1, 1], None]]

对于空间分区,输入流水线必须采用 tf.data 格式,并且 per_host_input_for_training 值 (train_batch_size) 必须设置为 PER_HOST_V2。您指定的 num_cores_per_replica 会确定您可以将图片拆分成的最大分区数。

input_partition_dims 值会提供一个包含两个元素的列表:feature_partition_dimslabel_partition_dims,用于描述如何对输入张量进行分区。feature_partition_dimslabel_partition_dims 的结构必须符合 input_fn 中的特征和标签结构。针对标签分区指定“无”,以便标签不会被拆分。

运行具有空间分区功能的参考模型

2D 对象检测

RetinaNet 是一种对象检测模型,可通过边界框定位图片中的对象并对所标识的对象进行分类。使用非常大的图片训练模型有一定的难度,因为适合单个 Cloud TPU 核心的最大图片(每台设备批量为 8 个)为 1280x1280。如下所示,您可以通过在 8 个 TPU 核心之间对模型进行空间分区来训练超过 4 倍的图片。

图片大小 TPU 类型 TPU 核心 全局批量大小 单步用时
1280x1280 v3-8 8 64 910 ms
2560x2560 v3-64 64 64 822 ms

3D 图片分割

3D UNet 是一种常见的密集型 3D 分割模型,已广泛应用于医学影像领域。CT 图片的原始分辨率最大可达 256x256x256,不适合单个 TPU 核心,因此研究人员通常必须对图片进行降采样。借助 TPU 空间分区,您可以使用 16 向空间分区直接放入原始分辨率图片。

图片大小 TPU 类型 TPU 核心 全局批量大小 单步用时
128x128x128 v3-8 8 32 3.428s
2560x2560 v3-64 64 32 3.02s

其他资源