tensorflow - 在 TPU 上使用大型 tensorflow 数据集

Question

以 TPU 可接受的方式缓存/生成数据集的最佳策略是什么？

到目前为止，我设法在自己创建的数据集上训练 tensorflow 模型。每个数据点都基于大型时间序列进行大量设计，使用基于 numpy、pandas、scipy 和其他 python 包构建的定制逻辑。最终的数据集创建步骤如下所示:

train_ds = tf.data.Dataset.from_generator(
    generator=data_gen,
    output_types=(tf.float32, tf.float32),
    output_shapes=([None, len(cols_of_interest)],()),
    args=([True, datafile, REGRESSION, num_features])
)

模型很好，在使用 CPU 时收敛。当我转向在 Google Colab 中使用 TPU 时，我收到与无法在 TPU 上运行我的 data_gen 函数相关的错误:

/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

NotFoundError:  No registered 'PyFunc' OpKernel for 'CPU' devices compatible with node {{node PyFunc}}
    .  Registered:  <no registered kernels>

     [[PyFunc]]
     [[IteratorGetNext]]
    Encountered when executing an operation using EagerExecutor. This error cancels all future operations and poisons their output tensors.
Additional GRPC error information:
{"created":"@1576029311.304590975","description":"Error received from peer","file":"external/grpc/src/core/lib/surface/call.cc","file_line":1039,"grpc_message":" No registered 'PyFunc' OpKernel for 'CPU' devices compatible with node {{node PyFunc}}\n\t.  Registered:  <no registered kernels>\n\n\t [[PyFunc]]\n\t [[IteratorGetNext]]\n\tEncountered when executing an operation using EagerExecutor. This error cancels all future operations and poisons their output tensors.","grpc_status":5} [Op:__inference_distributed_function_5899]

Function call stack:
distributed_function -> distributed_function

在我的策略中，我从一个单一的原始数据文件生成多个数据点。这一切都根据文件的内容而有所不同，所以我无法预先知道有多少数据点。文件数量和最终数据点数量之间没有相关性。我还查看了 TimeseriesGenerator，但它对我不起作用。

我正在处理医疗传感器数据，因此原始文件在几十 GB 的范围内。将所有内容保存在一个大文件中实际上行不通。

以 TPU 可接受的方式缓存/生成数据集的最佳策略是什么？

Answer 1

我从未使用过 TPU，但您的错误似乎与 CPU/TPU 混合执行有关。 TensorFlow 数据集 API 使用图形模式执行，因此要求您的操作可用于您的加速器，对于标识为 PyFunc 操作的“自定义 python 生成器”可能不是这种情况。这是我的理解，可能是错误的。

顺便说一句，解决此问题的一种方法是将数据集生成和 TPU 训练分开。

使用第一个任务生成数据集并将其保存到磁盘。 TFRecord 格式允许您将数据集拆分为许多小文件而不是一个大文件。您还可以寻找一些尺寸优化:compress你的数据，你可以使用 sparse tensors等

然后你就可以使用更经典的tf.data.TFRecordDataset了，我认为它比 GeneratorDataset 更适合 TPU 执行。