tf.data: "mixing" batch sizes?

Question

I am working in TF 1.4.1 using the Dataset API. I have two sets of input data, x and y, with the same number of elements each. For both x and y, I have large context arrays context_x and context_y, say of size 10000-by-10.

The context varies for each input batch (although it is the same for each sample within a batch), so it makes sense to link them together in the input pipeline. I cannot store all the context arrays for all of the batches in the graph and then read from there, since storing all of them statically in the graph is prohibitive in terms of the memory required. What I can do is input a single context array for the current batch, and I want to include this in my input pipeline.

Also note that in my graph the context array goes through some convolutional layers, which effectively reduce it to a manageable size, much smaller than original, so that I then tile it with other features specific to the samples in the batch and carry on with the rest of my task. So even though I will need to eventually replicate to batch size, I can do this on an feature vector extracted from the context array that is of much smaller size.

I am using the following type of code to build a dataset that should feed a batch of x and y samples along with their context to my graph:

import tensorflow as tf
import numpy as np

# Small data input
x = np.arange(100)
y = np.arange(100)

# Large context array for both x  and y
context_x = np.random.rand(1, 10000, 10)
context_y = np.random.rand(1, 10000, 10)

# Create datasets
dataset_x = tf.data.Dataset.from_tensor_slices(x)
dataset_y = tf.data.Dataset.from_tensor_slices(y)

# same context should be repeated for every data item
dataset_context_x = tf.data.Dataset.from_tensor_slices(context_x)
dataset_context_x = dataset_context_x.repeat()
dataset_context_y = tf.data.Dataset.from_tensor_slices(context_y)
dataset_context_y = dataset_context_y.repeat()

dataset = tf.data.Dataset.zip((dataset_x, dataset_context_x))
dataset = dataset.concatenate( tf.data.Dataset.zip((dataset_y, dataset_context_y)) )
dataset = dataset.batch(32)

iterator = dataset.make_initializable_iterator()
(x_iter, context_iter) = iterator.get_next()
with tf.Session() as sess:
    sess.run(iterator.initializer)
    while True:
        try:
            xi, ci = sess.run([x_iter, context_iter])
            print(xi.shape, ci.shape)
        except tf.errors.OutOfRangeError:
            break

The output indicates that the large context arrays are replicated for every sample x[i] and y[i]:

((32,), (32, 10000, 10))
((32,), (32, 10000, 10))
((32,), (32, 10000, 10))
((32,), (32, 10000, 10))
((32,), (32, 10000, 10))
((32,), (32, 10000, 10))
((8,), (8, 10000, 10))

This will waste a lot of memory, since all of the 32 10000-by-10 slices will be identical! How should I use the Dataset API so that this unnecessary copying of the context arrays is avoided, and I get an output such as ((32,), (1, 10000, 10)) for each batch? You could think of this as mixing batch sizes, 32 for x and y while 1 for the context arrays.

Answer 1

Ok, here's my tentative solution. Please note that I'm assuming your data is somehow ordered, so that when you build batches of x the next context_x you read is always the one related to the current batch.

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # running on CPU
import tensorflow as tf
import numpy as np

# Small data input
x = np.arange(100)
y = np.arange(100)

# Large context array for both x  and y
context_x = np.random.rand(1, 10000, 10)
context_y = np.random.rand(1, 10000, 10)

# Create datasets
dataset_x = tf.data.Dataset.from_tensor_slices(x).batch(32)
dataset_y = tf.data.Dataset.from_tensor_slices(y).batch(32)

# same context should be repeated for every data item
dataset_context_x = tf.data.Dataset.from_tensor_slices(context_x)
dataset_context_x = dataset_context_x.repeat() # here just for demonstration purposes. Ideally you'll have enough context data to match the batches
dataset_context_y = tf.data.Dataset.from_tensor_slices(context_y)
dataset_context_y = dataset_context_y.repeat() # here just for demonstration purposes. Ideally you'll have enough context data to match the batches

dataset = tf.data.Dataset.zip((dataset_x, dataset_context_x))
dataset = dataset.concatenate( tf.data.Dataset.zip((dataset_y, dataset_context_y)) ) # This stacks all 'x' samples on top of all 'y' samples. Is this really what you wanted?

iterator = dataset.make_initializable_iterator()
(x_iter, context_iter) = iterator.get_next()
with tf.Session() as sess:
    sess.run(iterator.initializer)
    while True:
        try:
            xi, ci = sess.run([x_iter, context_iter])
            print(xi.shape, ci.shape)
        except tf.errors.OutOfRangeError:
            break

In your implementation, remove the dataset_context_* = dataset_context_*.repeat() lines.

The key difference with your pipeline is that I'm batching x before zipping it with the context, so that the context doesn't get replicated. This, however, requires you to be careful in handling the data loading (hence my assumption above).