Tensorflow 2.0 数据集批处理无法正常工作

Question

Tensorflow 2.0 数据集 api 的批处理没有按我预期的那样工作。

我做了一个这样的数据集。

self.train_dataset = tf.data.Dataset.from_generator(generator=train_generator, output_types=(tf.float32, tf.float32), output_shapes=(tf.TensorShape([6]), tf.TensorShape([])))

这将产生 DatasetV1Adapter 形状:((6,), ())，类型:(tf.float32, tf.float32)，
我对这个数据集应用了 tf.data.Dataset 中的批处理函数。

self.train_dataset.batch(1024)

产生 DatasetV1Adapter 形状:((None, 6), (None,))，类型:(tf.float32, tf.float32)，改变批量大小根本没有帮助。

从批次的官方描述来看，

The components of the resulting element will have an additional outer dimension, which will be batch_size (or N % batch_size for the last element if batch_size does not divide the number of input elements N evenly and drop_remainder is False). If your program depends on the batches having the same outer dimension, you should set the drop_remainder argument to True to prevent the smaller batch from being produced.

我认为这个函数会起作用的方式是制作 [batch, 6], [batch,] 但效果不佳。

我最初使用 pytorch，最近开始使用 TF 2.0，需要一些关于正确批处理的帮助。提前致谢。

Answer 1

您可以通过设置获得所需的结果，

train_dataset = train_dataset.batch(2, drop_remainder=True)

drop_remainder=False默认情况下。在这种情况下，第一个维度 必须是 None因为(很可能)会有一批 < batch_size数据集末尾的元素，因为样本数不能被 batch_size整除.