Tensorflow - tf.data.Dataset.from_generator 예제

모델에 데이터를 공급하는 방식 중 하나인 tf.data.Dataset.from_generator를 살펴봅니다.

 

데이터가 너무 커서 메모리에 모두 로드할 수 없는 경우 데이터를 한번에 모두 읽어오는 대신에 디스크에서 데이터를 배치(batch) 단위로 로드하도록 할 수 있습니다. 

2021. 11. 19  최초작성

tf.data.Dataset.from_generator

 

우선 모델이 필요로 하는 데이터를 생성하는 generator 함수가 필요합니다. 

이 함수는 return 문 대신에 yield 문을 사용합니다. 

 

여기에서 데이터셋은 28 x 28 크기의 이미지와 10개의 클래스 중 하나임을 나타내는 정수 쌍으로 구성된 1000개의 데이터라고 가정합니다. 

generator 함수에서 다음처럼 가상의 데이터를 생성하여 yield 문을 호출합니다. 

 

이 generator 함수를 사용하여 TensorFlow dataset을 만들 수 있습니다. 

tf.data.Dataset.from_generator 함수의 아규먼트는 다음과 같습니다. 

 

output_shapes 아규먼트는 선택 사항이지만 output_types 아규먼트는 지정해야 합니다. output_types 아규먼트에서  generator 함수에서 리턴되는 값의 타입을 지정합니다. 첫 번째 리턴되는 값은 tf.float32 타입의 2차원 배열이고  두 번째 값은 tf.int16 타입의 1차원 배열이 됩니다.  

 

다음처럼 generator 함수를 사용하여  데이터셋의 데이터를 가져올 수 있습니다. 

import tensorflow as tf import numpy as np def our_generator():     for i in range(1000):       x = np.random.rand(28,28)       y = np.random.randint(1,10, size=1)       yield x,y dataset = tf.data.Dataset.from_generator(our_generator, (tf.float32, tf.int16)) for batch, (x,y) in enumerate(dataset):   pass print("batch: ", batch) print("Data shape: ", x.shape, y.shape) #batch:  999 #Data shape:  (28, 28) (1,)

batch

 

이전 예제에서는 for문에서 dataset 객체로부터 데이터를 하나씩 가져옵니다.

여기에서는 일정 크기의 데이터를 하나의 배치(batch)로 묶어서 가져옵니다.

import tensorflow as tf import numpy as np def our_generator():     for i in range(1000):       x = np.random.rand(28,28)       y = np.random.randint(1,10, size=1)       yield x,y dataset = tf.data.Dataset.from_generator(our_generator, (tf.float32, tf.int16)) dataset = dataset.batch(10) for batch, (x,y) in enumerate(dataset):   pass print("batch: ", batch) print("Data shape: ", x.shape, y.shape) # batch:  99 # Data shape:  (10, 28, 28) (10, 1)

repeat

 

repeat(count=2)를 사용하면 데이터셋을 2번 반복하여 사용합니다. 

데이터셋을 무한히 반복하려면 count=-1 아규먼트를 사용하면 됩니다. 

 

import tensorflow as tf import numpy as np def our_generator():     for i in range(1000):       x = np.random.rand(28,28)       y = np.random.randint(1,10, size=1)       yield x,y dataset = tf.data.Dataset.from_generator(our_generator, (tf.float32, tf.int16)) dataset = dataset.batch(10) dataset = dataset.repeat(count=2) for batch, (x,y) in enumerate(dataset):   pass print("batch: ", batch) print("Data shape: ", x.shape, y.shape) #batch:  199 #Data shape:  (10, 28, 28) (10, 1)

shuffle

 

Epoch 마다 데이터셋을 뒤섞기위해 shuffle를 사용합니다. 

 

import tensorflow as tf import numpy as np def our_generator():     for i in range(1000):       x = np.random.rand(28,28)       y = np.random.randint(1,10, size=1)       yield x,y dataset = tf.data.Dataset.from_generator(our_generator, (tf.float32, tf.int16)) dataset = dataset.batch(10) dataset = dataset.shuffle(buffer_size=1000) for batch, (x,y) in enumerate(dataset):   pass print("batch: ", batch) print("Data shape: ", x.shape, y.shape) #batch:  99 #Data shape:  (10, 28, 28) (10, 1)

 

데이터셋에서 가져온 데이터로 buffer_size 크기의 버퍼에 채운 다음 이 버퍼에서 이 요소를 무작위로 샘플링합니다.

buffer_size>=dataset_size 만족하도록 해야 만족스러운 무작위 샘플링이 됩니다.

참고

 

https://vak.ai/TensorFlow2.0-dataset/