Tensorflow:tf.data.Dataset,不能在组件0中批处理具有不同形状的张量

Question

我的输入管道中出现以下错误:

tensorflow.python.framework.errors_impl.InvalidArgumentError:无法在组件0中批量生成具有不同形状的张量.第一个元素具有形状[2,48,48,3],元素1具有形状[27,48,48,3].

用这个代码

dataset = tf.data.Dataset.from_generator(generator,
                                         (tf.float32, tf.int64, tf.int64, tf.float32, tf.int64, tf.float32))

dataset = dataset.batch(max_buffer_size)

这是完全合乎逻辑的,因为批处理方法尝试创建(batch_size,？,48,48,3)Tensor.但是我希望它为这种情况创建一个[29,48,48,3] Tensor.所以连接而不是堆栈.这可能与tf.data有关吗？

我可以在生成器函数中用Python进行连接,但我想知道这是否也可以使用tf.data管道

Answer 1

第一种情况:我们希望输出具有固定的批量大小

在这种情况下,生成器生成形状值,[None, 48, 48, 3]其中第一维可以是任何形状.我们想批量处理这个输出[batch_size, 48, 48, 3].如果我们直接使用tf.data.Dataset.batch,我们会有错误,所以我们需要先解开.

为此,我们可以tf.contrib.data.unbatch在批处理之前使用这样的方法:

dataset = dataset.apply(tf.contrib.data.unbatch())
dataset = dataset.batch(batch_size)

下面是一个完整的例子,其中发电机的产率[1],[2, 2],[3, 3, 3]和[4, 4, 4, 4].

我们无法直接批量处理这些输出值,因此我们将其解包然后批处理:

def gen():
    for i in range(1, 5):
        yield [i] * i

# Create dataset from generator
# The output shape is variable: (None,)
dataset = tf.data.Dataset.from_generator(gen, tf.int64, tf.TensorShape([None]))

# The issue here is that we want to batch the data
dataset = dataset.apply(tf.contrib.data.unbatch())
dataset = dataset.batch(2)

# Create iterator from dataset
iterator = dataset.make_one_shot_iterator()
x = iterator.get_next()  # shape (None,)

sess = tf.Session()
for i in range(5):
    print(sess.run(x))

这将打印以下输出:

[1 2]
[2 3]
[3 3]
[4 4]
[4 4]

第二种情况:我们想要连接可变大小的批次

更新(03/30/2018):我删除了之前使用分片的答案,这会大大降低性能(请参阅注释).

在这种情况下,我们希望连接固定数量的批次.问题是这些批次的大小可变.例如,数据集产生[1]并且[2, 2]我们希望得到[1, 2, 2]输出.

这里解决这个问题的一个快速方法是创建一个围绕原始生成器的新生成器.新生成器将生成批量数据.(感谢纪尧姆的想法)

下面是一个完整的例子,其中发电机的产率[1],[2, 2],[3, 3, 3]和[4, 4, 4, 4].

def gen():
    for i in range(1, 5):
        yield [i] * i

def get_batch_gen(gen, batch_size=2):
    def batch_gen():
        buff = []
        for i, x in enumerate(gen()):
            if i % batch_size == 0 and buff:
                yield np.concatenate(buff, axis=0)
                buff = []
            buff += [x]

        if buff:
            yield np.concatenate(buff, axis=0)

    return batch_gen

# Create dataset from generator
batch_size = 2
dataset = tf.data.Dataset.from_generator(get_batch_gen(gen, batch_size),
                                         tf.int64, tf.TensorShape([None]))

# Create iterator from dataset
iterator = dataset.make_one_shot_iterator()
x = iterator.get_next()  # shape (None,)


with tf.Session() as sess:
    for i in range(2):
        print(sess.run(x))

这将打印以下输出:

[1 2 2]
[3 3 3 4 4 4 4]