使用 from_generator 创建的 TensorFlow 数据集未按 batch_size 迭代

【问题标题】：Tensorflow Dataset created with from_generator not iterating by batch_size使用 from_generator 创建的 TensorFlow 数据集未按 batch_size 迭代
【发布时间】：2020-03-27 19:14:09
【问题描述】：

我从生成器创建了一个 tensorflow 数据集，但不知道如何通过 batch_size 对其进行迭代

def ds_gen():
    x = np.random.random((10, 10, 3))
    y = np.random.random((2))
    yield x, y

def create_tf_dataset():
    dataset = tf.data.Dataset.from_generator(ds_gen, output_types=(tf.float32, tf.float32), output_shapes=((10, 10, 3), (2,)))
    return dataset

ds = create_tf_dataset()
ds = ds.batch(10)
for x_batch, y_batch in ds:
  print(x_batch.shape, y_batch.shape)

此代码不断循环批处理大小为 1 而不是 10

【问题讨论】：

标签： tensorflow tensorflow-datasets

【解决方案1】：

请参考下面的代码以按批量大小进行迭代

import numpy as np
import tensorflow as tf 

def ds():
    for i in range(1000):
      x = np.random.rand(10,10,3)
      y = np.random.rand(2)
      yield x,y

ds = tf.data.Dataset.from_generator(ds, output_types=(tf.float32, tf.float32), output_shapes=((10, 10, 3), (2,)))

ds = ds.batch(10)

for batch, (x,y) in enumerate(ds):
  pass
print("Data shape: ", x.shape, y.shape)

输出：

Data shape:  (10, 10, 10, 3) (10, 2)

如果您更改ds = ds.batch(1)，则输出将为Data shape: (1, 10, 10, 3) (1, 2)

【讨论】：