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

Question

我的输入管道中有以下错误：

tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot batch tensors with different shapes in component 0. First element had shape [2,48,48,3] and element 1 had shape [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)

当批处理方法尝试创建（批处理大小？，48，48，3）张量。但是我希望它为这个例子创建一个[29，48，48，3]张量。所以连接而不是堆栈。用tf.data可以吗？

我可以在generator函数中使用Python进行连接，但我想知道tf.data管道是否也可以这样做

Answer 1

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

在这种情况下，生成器生成形状[None, 48, 48, 3]的值，其中第一个维度可以是任何内容。我们希望对其进行批处理，以便输出为[batch_size, 48, 48, 3]。如果我们直接使用tf.data.Dataset.batch，我们将有一个错误，因此我们需要首先取消批处理。

为此，我们可以在批处理之前使用^{}方式：

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]

---

第二种情况：我们希望连接可变大小的批

更新（2018年3月30日）：我删除了之前使用分片的答案，这会大大降低性能（请参阅注释）。

在这种情况下，我们希望连接固定数量的批。问题是这些批次的大小不一。例如，数据集产生[1]和[2, 2]，我们希望得到[1, 2, 2]作为输出。

解决这个问题的一个快速方法是在原来的生成器周围创建一个新的生成器。新的生成器将生成批处理数据。（感谢Guillaume的想法）

---

这里是一个完整的例子，生成器生成[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]