Bug fix: avoid concatenating TensorFlow Datasets linearly
@@ -95,13 +95,17 @@ All windows have the same number of elements.
def concat_datasets(datasets):
- concat_dataset = None
- for scene_dataset in datasets:
- if concat_dataset is None:
- concat_dataset = scene_dataset
+ """ Creates a TensorFlow Dataset from the given ordered list of TensorFlow Datasets """
+ # n-1 concatenations will lead to overflow (RecursionError) if list is too long
+ # Instead, concatenate them hierarchically
+ assert len(datasets) > 0
+ if len(datasets) == 1:
+ return datasets[0]
else:
- concat_dataset = tf.data.Dataset.concatenate(concat_dataset, scene_dataset)
- return concat_dataset
+ center = len(datasets) // 2
+ dataset_1 = concat_datasets(datasets[:center])
+ dataset_2 = concat_datasets(datasets[center:])
+ return tf.data.Dataset.concatenate(dataset_1, dataset_2)
def _example_count(length, frames, inner_stride, outer_stride):
