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

Question

I'm trying to use TF Datasets with a @tf.function to perform some preprocessing on a directory of images. Inside the tf function the image file is read as a RAW string tensor and I'm trying to take a slice from that tensor. The slice, the first 13 characters, represent info about .ppm images (header). I get an error: ValueError: Shape must be rank 1 but is rank 0 for 'Slice' (op: 'Slice') with input shapes: [], [1], [1] . Initially I was trying to directly slice the .numpy() attribute of the tensor ( filepath input parameter to the tf function), but I think it is semantically wrong to do this inside a tf function. It also didn't work as the filepath input tensor does not have a numpy() attribute (I don't understand why??). Outside of the tf function, eg in a jupyter notebook cell, I can iterate over the dataset and get individual items which have a numpy attribute and do a slice and all subsequent processing on it just fine. I do realize there may be a gap in my understanding of how TF works (I am using TF 2.0), so I hope someone can clarify what I missed in my readings. The purpose of the tf function is convert the ppm images to png, so there is a side effect of this function, but I did not get that far to find out if this is possible to do.

Here's the code:

@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

To manipulate string values in TF, have a look at the tf.strings namespace .

In this case, you can use 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 only operates on the Tensor objects, and doesn't work on their elements. Here, ppm_bytes is a scalar Tensor containing a single element of type tf.string , and whose value is the entire string contents of the file. So when you call tf.slice , it only looks at the scalar bit, and is not smart enough to realize that you actually want to take a slice of that element instead.