tf.dataset不追加批处理

时间:2018-08-13 07:42:35

标签: python tensorflow tensorflow-datasets

我想让tf.dataset工作。下面的代码示例正在运行,但是由于我使用了.batch(30),所以我希望输出的格式为(30,300,300,1)?

import tensorflow as tf
import numpy as np

input_array = np.random.normal(size=(300, 300, 3))

def own_generator():
    yield (input_array, input_array)

dataset = tf.data.Dataset.from_generator(own_generator, (tf.float32, tf.float32)).batch(30)
data_iter = dataset.make_initializable_iterator()

sess = tf.Session()
sess.run(data_iter.initializer)
test_arr = sess.run(data_iter.get_next())

for tuple_elemnt in test_arr:
    print(tuple_elemnt.shape)

输出为:

(1, 300, 300, 3)
(1, 300, 300, 3)

1 个答案:

答案 0 :(得分:1)

生成器编程错误。这是工作示例:

import tensorflow as tf
import numpy as np

input_array = np.random.normal(size=(300, 300, 3))

def own_generator():
    while True:
        yield input_array

dataset = tf.data.Dataset.from_generator(own_generator, tf.float32).batch(30)
data_iter = dataset.make_initializable_iterator()

sess = tf.Session()
sess.run(data_iter.initializer)
test_arr = sess.run(data_iter.get_next())

print(test_arr.shape)