时间戳

首页 > 解决方案 > TPU 上的 RaggedTensor

python - TPU 上的 RaggedTensor

问题描述

我正在尝试使用将RaggedTensortf.keras.layers.Input作为输入（ ）作为输入的 TensorFlow 训练神经网络。它在 CPU 和 GPU 上运行良好，但我真的很努力让它在 TPU 上运行。我想知道你们中的一些人是否设法使它工作（不一定要寻找直接的解决方案，尽管它会很棒，一些工具提示已经很棒了！）。到目前为止，错误消息对我来说已经足够明确了，但我现在还在为如何走得更远而苦苦挣扎。

到目前为止我做了什么：

  1. 我tf.data.Dataset用来从 TF_Records 读取数据，但我需要将其显式转换为DistributedDataset以禁用 prefecting。
strategy.experimental_distribute_dataset(
    dataset,
    tf.distribute.InputOptions(
        experimental_prefetch_to_device=False
    )
)
  1. 通过允许软设备放置，我得到了Compilation failure: Detected unsupported operations when trying to compile graph ... on XLA_TPU_JIT: RaggedTensorToTensor可以（某种程度上）修复的问题：
tf.config.set_soft_device_placement(True)
  1. 我现在陷入困境Compilation failure: Input 1 to node '.../RaggedReduceSum/RaggedReduce/RaggedSplitsToSegmentIds/Repeat/SequenceMask/Range' with op Range must be a compile-time constant.。我完全理解为什么会出现此错误，我完全了解TPU 上的可用操作，特别是大多数动态操作应在编译时确定以在 TPU 上运行。但我想不出我怎么能用 TPU 来使用那些参差不齐的张量......

任何想法将不胜感激:)

PS：自从 2020 年 7 月的这个答案以来，我还没有看到 TensorFlow 团队关于 TPU 上 RaggedTensors 的太多消息，但我可能已经错过了很多关于它的消息……如果我指向 git 线程对我来说已经很棒了可以调查更多。

标签： pythontensorflowtpuragged-tensors

解决方案

推荐阅读