대용량 훈련 데이터 처리 - Generator로 TF Dataset 만들기

1. 너무 큰 데이터

  tf.data.Dataset는 Tensorflow의 훈련 데이터를 다룰 때 참 편리합니다. Padding, Batch, Shuffle, Map 등을 다 지원해주니까요. 일단 데이터를 Dataset으로 변환만 시키면 그다음부터는 아주 편리하게 사용할 수 있습니다. from_tensor_slices를 쓰면 numpy array를 바로 Dataset으로 변환해주니까 더할나위 없이 편리하고요.

  하지만! 이 방법에는 치명적인 문제가 있으니... 바로 변환하려는 전체 데이터가 메모리에 올릴 수 있는 크기여야 한다는 겁니다. 물론 가상 메모리를 써서 메모리를 가짜로 늘릴 수는 있지만 이 순간부터 Thrashing으로 속도가 급격하게 떨어지고, Out-of-memory로 죽기도 합니다. 이 문제가 굉장히 현실적인 이유는 실제 환경에서 우리가 다루는 훈련 데이터는 대부분 RAM에 올릴 수 없을 정도로 크기 때문입니다. 자연어 처리든, 이미지 처리든 마찬가지고요.

  이럴 때 아주 요긴하게 쓸 수 있는 방법이 바로 Dataset의 from_generator입니다. from_generator를 사용하면 한번에 데이터를 메모리에 다 로딩하는 것이 아니고, 필요할 때만 python generator를 통해서 데이터를 가져옵니다. python generator의 목적 자체가 무한한 데이터를 다루거나, 데이터를 필요할 때 가져오는 lazy loading을 위한 것이니 딱 맞아떨어지는 셈이죠.

  그래서 from_tensor_slices를 실행하면 numpy array를 변환하기 위해서 엄청나게 오랜 시간이 걸리곤 하지만, from_generator를 호출하면 바로 실행이 끝납니다. 실제로는 데이터를 로딩하지 않으니까요.

게으른 것이 다 나쁜 건 아닙니다.

2. from_generator 

from_generator(
    generator, output_types, output_shapes=None, args=None
)

  from_generator의 첫번째 인자는 데이터를 제공해 줄 generator입니다. 여기에서 중요한 점은 첫 번째 인자는 생성된 generator가 아니고, 호출하면 generator를 돌려주는 함수(또는 callable)여야한다는 점입니다. 말이 굉장히 난해한데, 역시 이럴 때는 코드가 답입니다.

def gen():
    for i in range(10):
        yield i
        
my_generator = gen()
print(type(my_generator))
print(type(gen))

<class 'generator'>
<class 'function'>

  gen 함수를 실행했을 때 반환된 값은 generator class이고 gen 자체는 function class입니다. 즉, from_generator 에게는 my_generator가 아닌, gen을 인자로 전달해야 합니다. 이유는 epoch 때문입니다. generator가 모든 데이터를 다 소진하고 StopIteration을 던지면, Tensorflow는 한 epoch을 다 돌았다고 생각합니다. 그럼 다음 epoch을 시작해야 하는데 이미 generator는 데이터를 다 소진한 상태이기 때문에 처음부터 데이터를 다시 가져 올 방법이 없습니다. generator는 한 방향으로만 진행하기 때문이죠.

일방통행이에요!

  하지만 generator를 생성할 수 있는 함수(또는 callable)를 인자로 주면 이 문제는 해결됩니다. generator를 새로 만들어서 시작하면 되니까요.

 

  두번째 인자는 generator가 돌려주는 데이터의 Type이고, 세 번째 인자는 generator 돌려주는 데이터의 Shape입니다. 나중에 나올 예시 코드에서 실제 예를 보도록 하죠.

  네 번째 인자는 generator에게 전달해 줄 인자들입니다. Tuple이고요. generator가 인자를 취하지 않는다면 무시해도 되지만 인자가 필요하다면 args를 통해 넘겨주면 됩니다. 한 가지 주의 사항은 numpy array 형태로 넘어가기 때문에 모든 python 타입을 다 넣을 수는 없습니다. 기본적인 수치, 문자열은 가능하고요.

 

  성공적으로 모든 값을 주고 from_generator가 호출되면 Dataset을 반환합니다. 그다음부터는 일반적으로 Dataset을 쓰는 것과 같고요.

3. 실전 코드!

  이렇게 설명만 보면 이해가 좀 어려우니 iris dataset을 가지고 간단히 예시 코드를 만들어 보겠습니다. iris dataset은 150개 examples 밖에 안 되지만 어마어마하게 커서 메모리에 다 넣기가 힘들다고 가정을 해보겠습니다.

 

  * iris dataset은 꽃의 특징을 나타내는 숫자를 가지고 세 가지 붓꽃 중 하나로 분류하는 classification dataset입니다.

5.1,3.5,1.4,0.2,Iris-setosa
4.9,3.0,1.4,0.2,Iris-setosa
4.7,3.2,1.3,0.2,Iris-setosa
4.6,3.1,1.5,0.2,Iris-setosa
5.0,3.6,1.4,0.2,Iris-setosa
5.4,3.9,1.7,0.4,Iris-setosa
.
.
.
.

  앞에 4개의 숫자가 각 꽃의 특징이고, 마지막 열이 붓꽃의 종류입니다.

  먼저 generator 함수를 만들어 보겠습니다.

def get_iris_data(path):
    label_to_index = {
        'Iris-setosa': 0,
        'Iris-versicolor': 1,
        'Iris-virginica': 2
    }
    
    for line in open(path):
        tokens = line.strip().split(',')
        features = [float(token) for token in tokens[:4]]
        label = label_to_index[tokens[4]]
        
        yield (features, label)

  전달받은 경로에서 한 줄씩 읽으면서 ,를 경계로 앞으로 4가지는 float로 변환해서 features라는 이름으로 묶어주고, 마지막 열은 숫자 값으로 변환해서 돌려줍니다.

gen = get_iris_data('./iris.data')
print(next(gen))

  이렇게 해보시면 첫 번째 데이터를 받아보실 수 있겠죠.

 

  다음으로는 위의 generator를 이용해서 Dataset을 만들어줍니다.

import tensorflow as tf

dataset = tf.data.Dataset.from_generator(get_iris_data,
                                         (tf.float64, tf.int64),
                                         (tf.TensorShape([None]), tf.TensorShape([])),
                                         args=('./iris.data',))

  첫번째 인자로 get_iris_data 함수(get_iris_data()이 값이 아님)를 전달합니다.

  두 번째 인자로는 generator가 반환하는 데이터의 Type을 전달하는데, 이 경우 features는 float, label은 int이기 때문에 이렇게 지정합니다.

  세 번째는 shape인데요, features는 1차원, label은 scalar 이기 때문에 위와 같이 지정합니다. features의 원소가 4개라는 것을 명시하고 싶으시다면 (tf.TensorShape([4]), tf.TensorShape([]))으로 하셔도 됩니다.

  마지막 인자는 get_iris_data에게 전달해 줄 path 인자이고요.

  

  이렇게 Dataset이 만들어지고 나면 그다음부터는 shuffle, batch 등을 그대로 사용할 수 있습니다.

dataset = dataset.shuffle(150).batch(8)

model = tf.keras.Sequential([
    tf.keras.layers.Dense(16, input_shape=(4,)),
    tf.keras.layers.Dense(4, activation='softmax')
    ])

model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy',])
model.fit(dataset, epochs=10)

4. 마무리

  pypthon generator와 Tensorflow의 from_generator를 사용해서 어마어마하게 큰 데이터를 훈련시킬 수 있는 방법을 알아보았습니다. 이 방법은 메모리를 아낄 수 있는 효과도 있지만, 모든 데이터를 한번에 메모리에 올리지 않아도 되기 때문에 훈련 시작을 빨리할 수 있다는 장점도 있습니다. 유사한 방법으로 tf.keras.utils.Sequence도 있는데요. 이 방법도 다음에 소개하도록 하겠습니다.

  그럼 Deep Learning 하겠다고 사서 게임만 돌리고 있는 그래픽 카드의 먼지를 털고 다시 Tensorflow의 세계로!