TensorFlow: Example implementation of a simple custom transformation_func for Dataset's apply method

user_3595278 0
Viewed 412

I am trying to implement a simple custom transformation_func for the apply method in the Dataset API, but didn't find the docs particularly helpful.

Specifically, my dataset contains video frames and corresponding labels: {[frame_0, label_0], [frame_1, label_1], [frame_2, label_2],...}.

I'd like to transform it so that it additionally contains the previous frame for each label: {[frame_0, frame_1, label_1], [frame_1, frame_2, label_2], [frame_2, frame_3, label_3],...}.

This could probably be achieved by doing something like tf.data.Dataset.zip(dataset, dataset.skip(1)), but then I would have duplicated labels.

I have not been able to find a reference implementation of a transformation_func. Is anyone able to get me started on this?

tensorflow-datasets tensorflow python
1 Answers

apply is simply a convenience to use with existing transformation functions, ds.apply(func) is pretty much the same as func(ds), only in a more "chainable" way. Here is one possible way to do what you want:

import tensorflow as tf

frames = tf.constant([  1,   2,   3,   4,   5,   6], dtype=tf.int32)
labels = tf.constant(['a', 'b', 'c', 'd', 'e', 'f'], dtype=tf.string)
# Create dataset
ds = tf.data.Dataset.from_tensor_slices((frames, labels))
# Zip it with itself but skipping the first one
ds = tf.data.Dataset.zip((ds, ds.skip(1)))
# Make desired output structure
ds = ds.map(lambda fl1, fl2: (fl1[0], fl2[0], fl2[1]))
# Iterate
it = ds.make_one_shot_iterator()
elem = it.get_next()
# Test
with tf.Session() as sess:
    while True:
        try: print(sess.run(elem))
        except tf.errors.OutOfRangeError: break

Output:

(1, 2, b'b')
(2, 3, b'c')
(3, 4, b'd')
(4, 5, b'e')
(5, 6, b'f')
user_17827920
Related