如何正确使用 tf.function 与 TensorFlow 数据集

Question

我正在尝试使用带有 @tf.function 的 TF 数据集对图像目录执行一些预处理。在 tf 函数中，图像文件被读取为 RAW 字符串张量，我试图从该张量中取出一个切片。切片（前 13 个字符）表示有关 .ppm 图像（标题）的信息。我收到一个错误：ValueError: Shape must be rank 1 but is rank 0 for 'Slice' (op: 'Slice') with input shapes: [], [1], [1]。最初我试图直接切片张量的 .numpy() 属性（filepath tf 函数的输入参数），但我认为在 内执行此操作在语义上是错误的tf 函数。它也不起作用，因为filepath 输入张量没有 numpy() 属性（我不明白为什么？？）。在 tf 函数之外，例如在 jupyter 笔记本单元格中，我可以遍历数据集并获取具有 numpy 属性的单个项目，并对其进行切片和所有后续处理。我确实意识到我对 TF 如何工作的理解可能存在差距（我使用的是 TF 2.0），所以我希望有人能澄清我在阅读中遗漏的内容。 tf 函数的目的是将 ppm 图像转换为 png，所以这个函数有一个副作用，但我并没有深入了解是否可以这样做。

代码如下：

@tf.function
def ppm_to_png(filepath):
    ppm_bytes = tf.io.read_file(filepath) #.numpy()
    bytes_header = tf.slice(ppm_bytes, [0], [13])
    # bytes_header = ppm_bytes[:13].eval()  # this did not work either with similar error msg
    .
    .
    .

import glob

files = glob.glob(os.path.join(data_dir, '00000/*.ppm'))
dataset = tf.data.Dataset.from_tensor_slices(files)
png_filepaths = dataset.map(ppm_to_png, num_parallel_calls=tf.data.experimental.AUTOTUNE)

Answer 1

要在 TF 中操作字符串值，请查看 tf.strings namespace。

在这种情况下，您可以使用tf.strings.substr：

@tf.function
def ppm_to_png(filepath):
  ppm_bytes = tf.io.read_file(filepath)
  bytes_header = tf.strings.substr(ppm_bytes, 0, 13)
  tf.print(bytes_header)

tf.slice 只对张量对象起作用，对它们的元素不起作用。这里，ppm_bytes 是一个标量张量，包含一个 tf.string 类型的元素，其值是文件的整个字符串内容。因此，当您调用 tf.slice 时，它只查看标量位，并且不够聪明，无法意识到您实际上想要取该元素的一部分。