Unable to get tf.data to match keras Sequence performance

user_76900 0
Viewed 65

I have a dataset consisting of an array of time series episodes. I have previously trained my neural network using a straightforward tf.keras.utils.Sequence, which windows each episode into equal length inputs with a definable shift. For example, data might look like this:

[
  [1, 2, 3, 4, 5, 6, 7], 
  [11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
]

And the inputs after applying a sliding window would be like this:

[1,  2,  3,  4,  5 ],
[2,  3,  4,  5,  6 ],
[3,  4,  5,  6,  7 ],
[11, 12, 13, 14, 15],
[12, 13, 14, 15, 16],
...
[16, 17, 18, 19, 20]

So we get a series of fixed length input vectors, ready to be batched.

Each episode is stored in a file and loaded into memory using np.load. The Keras Sequence code is straight-forward, like this:

def get_segment(self, segment_locator):
    filename = segment_locator[0]
    wav_data = self.cache[filename]
    r = wav_data[segment_locator[1]:segment_locator[2]]
    assert len(r) == self.fragment_length, "For %s, expected segment length %d, got %d" % (segment_locator, self.fragment_length, len(r))
    return r

def __getitem__(self, index):
    batch_set = self.segment_list[index * self.batch_size:(index + 1) * self.batch_size]
    batch_x = np.array([self.get_segment(segment_locator) for segment_locator in batch_set])
    batch_y = batch_x
    return batch_x, batch_y

Basically we have a "segment list" which is a precalculated list of segments identified by (filename, start, end) tuples. (End is in this case redundant since it will always be start + segment_length). We shuffle the list of segment selections, batch it and send that to model.fit.

I'm looking to switch to a multi-GPU setup and the recommendation I've read online is to use tf.data. I made a simple wrapper to convert my sequence to a tf dataset. Note that input == target in this example, so we zip together two copies of the same data sequence in the final step.

def straight_chunk_generator(self):
    for filename in self.filenames:
        for start_at in range(0, self.cache[filename].shape[0], self.shift):
            end_at = start_at + self.fragment_length
            if end_at > self.cache[filename].shape[0]:
                break
            yield self.cache[filename][start_at:end_at]

def as_tf_data(self):
    x_ds = tf.data.Dataset.from_generator(self.straight_chunk_generator, output_types=tf.float32)
    return tf.data.Dataset.zip((x_ds, x_ds))

The problem is that when I use this Dataset, epochs run about 5X as slowly. What's happening here? prefetch and cache seem to have no effect.

I also tried a more "pure" tf.data variation:

def as_tf_data0(self):
    dataset = tf.data.Dataset.from_tensor_slices(tf.ragged.constant([self.cache[filename] for filename in self.filenames]))

    def window_subseries(subarray):
        return tf.data.Dataset.from_tensor_slices(subarray).window(self.fragment_length, shift=self.shift, drop_remainder=True)

    x_ds = dataset.flat_map(window_subseries).flat_map(lambda x: x.batch(self.fragment_length))
    return tf.data.Dataset.zip((x_ds, x_ds))

(This got oddly complicated because window doesn't return the original data shape, but a series of datasets so they have to be "unwrapped", which x.batch does. In theory.)

This version is also inexplicably slow.

I am having difficulty understanding why my code is running slowly. I understand that the tf.data.Dataset.from_generator function can be slow due to the Python GIL, but I don't understand why it would be so much slower than the pure Python Sequence. And shouldn't adding a .cache() to the pipeline effectively force ahead of time evaluation of that part anyhow?

If the generator is the problem, why didn't my "native" implementation using from_tensor_slices, flat_map etc improve things? I thought that would lead to some kind of C++ optimised execution.

tensorflow-datasets keras tensorflow python
1 Answers

I ended up coming up with a solution for single GPU systems. The following code works well and is much faster than Sequence in my testing:

# index_list contains the offsets each segment as int32s.
host_array = np.concatenate([self.cache[filename] for filename in self.filenames])
with tf.device('/gpu:0'):
    device_array = tf.constant(host_array)
    frag_length = tf.constant([self.fragment_length])
    ds = tf.data.Dataset.from_tensor_slices(index_list)
    @tf.function
    def get_segment_tuple(index):
        sample = tf.slice(device_array, [index], frag_length)
        return (sample, sample)

    if self.shuffle:
        ds = ds.shuffle(len(index_list))

    x_ds = ds.map(get_segment_tuple, num_parallel_calls=tf.data.AUTOTUNE)
    return x_ds.batch(self.batch_size)

The techniques employed are:

  • ensure the backing vector is in GPU memory
  • since input=target in this example, call tf.slice just once and use the same slice twice (instead of zipping together two instances of the same dataset)
  • try to avoid giving any reason for TF to copy data from GPU-CPU (DtoH)
  • ...for example by removing that zip call just in case it lead to copying
  • ...and placing map after shuffling, maybe shuffling doesn't have a GPU kernel
  • place the segment length in an on device constant

I am not entirely satisfied with my understanding of this problem, but this solution works and the input time per step is now less than 0.1%. (I believe this to be a downside of Tensorflow in that it tries to do everything automatically. This makes it more difficult to understand what will and will not work-- and why.)

One flaw with this solution is that I am unsure of how to use it with multiple GPUs. Each one would need its own device_array in local memory and be able to know which one to use.

user_769000
Related