텐서플로 Dataset

import tensorflow as tf

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6, 7, 8])

for batch in dataset:
    print(batch)

tf.Tensor(1, shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)
tf.Tensor(3, shape=(), dtype=int32)
tf.Tensor(4, shape=(), dtype=int32)
tf.Tensor(5, shape=(), dtype=int32)
tf.Tensor(6, shape=(), dtype=int32)
tf.Tensor(7, shape=(), dtype=int32)
tf.Tensor(8, shape=(), dtype=int32)

for batch in dataset.batch(2):
    print(batch)

tf.Tensor([1 2], shape=(2,), dtype=int32)
tf.Tensor([3 4], shape=(2,), dtype=int32)
tf.Tensor([5 6], shape=(2,), dtype=int32)
tf.Tensor([7 8], shape=(2,), dtype=int32)

for batch in dataset.batch(4):
    print(batch)

tf.Tensor([1 2 3 4], shape=(4,), dtype=int32)
tf.Tensor([5 6 7 8], shape=(4,), dtype=int32)

def gen():
    for i in range(10):
        yield [i] * i

dataset = tf.data.Dataset.from_generator(gen, output_types=tf.int32)

for batch in dataset.padded_batch(4, padded_shapes=(None, )):
    print(batch)

tf.Tensor(
[[0 0 0]
 [1 0 0]
 [2 2 0]
 [3 3 3]], shape=(4, 3), dtype=int32)
tf.Tensor(
[[4 4 4 4 0 0 0]
 [5 5 5 5 5 0 0]
 [6 6 6 6 6 6 0]
 [7 7 7 7 7 7 7]], shape=(4, 7), dtype=int32)
tf.Tensor(
[[8 8 8 8 8 8 8 8 0]
 [9 9 9 9 9 9 9 9 9]], shape=(2, 9), dtype=int32)