Tensorflow에서 커스텀 데이터로더 만들기(Custom Dataloader, Sequence)
Sequence를 사용하여 Custom Dataset 만들기
pytorch에서 보통 데이터 로드할때 torch.utils.data.DataSet을 많이 사용하는데요,
tensorflow 2.0 이상 버전에서도 비슷하게 custom dataset loader를 만드는 방법이 있습니다.
tensorflow.keras.utils.Sequence를 사용하는 건데요, 사용하는 방법을 알아보겠습니다.
우선 __init__, __len__, __getitem__ 함수 세 가지는 필수적으로 만들어 주어야 합니다.
초기화 함수(__init__)
우선 아래와 같이 Sequence모듈을 import하고 Sequence를 상속받는 클래스를 생성해줍니다.
다른 클래스와 마찬가지로 __init__ 함수를 사용하여 필요한 데이터(x, y)를 받아옵니다.
from tensorflow.keras.utils import Sequence
class CustomDataloader(Sequence):
def __init__(self, x_set, y_set, batch_size, shuffle=False):
self.x, self.y = x_set, y_set
self.batch_size = batch_size
self.shuffle=shuffle길이 함수(__len__)
데이터 로더의 전체 길이를 반환하는 함수입니다.
예를 들어 50000개의 데이터가 있는데, batch size가 100이라면 loader의 길이는 500이 됩니다.
하지만 batch size가 딱 맞아 떨어지지 않는 경우도 있는데요, 그럴 땐 올림을 해줍니다.
다시 예를 들어, 49990개의 데이터가 있고, batch size가 100이라면 loader의 길이는 위와 마찬가지고 500이 됩니다.
배치 단위로 데이터를 가져올 텐데 이렇게 하면 마지막 배치는 90개가 되어 버리는 데이터 없이 모두 사용이 가능합니다. 이 때, 올림을 위해 math 모듈을 사용하겠습니다.
import math
def __len__(self):
return math.ceil(len(self.x) / self.batch_size)index값에 따라 데이터를 반환하는 함수(__getitem__)
이 함수가 중요한데요, 위에서 설명했듯이 10개의 데이터가 있고 배치 사이즈가 2개인 경우 총 길이는 5입니다.
이때 index가 0~4에 따라 배치를 반환 하는데요, 소스를 통해 알아보겠습니다.
def __getitem__(self, idx):
indices = self.indices[idx*self.batch_size:(idx+1)*self.batch_size]
batch_x = [self.x[i] for i in indices]
batch_y = [self.y[i] for i in indices]
return np.array(batch_x), np.array(batch_y)여기서 self.indices는 전체 데이터의 index 배열 입니다. (아래 on_epoch_end함수 참고) 데이터가 10개인 경우 각각 0~9번 째의 데이터 index 배열을 의미합니다. 그 중 indices는 __getitem__함수를 호출한 idx에 맞는 배치사이즈 만큼의 index 배열입니다.
전체 데이터(self.indices)가 [0 1 2 3 4 5 6 7 8 9] 이고, batch_size가 2라고 한다면,
__getitem__(0)을 호출하면 indices는 [0, 1]이 되어 각 index에 맞는 데이터 [0, 1]를 반환할 것이고
__getitem__(1)을 호출하면 indices는 [2, 3]이 되고 각 index에 맞는 데이터 [2, 3]을 반환할 것입니다.
함수를 호출하는 idx에 따라 배치 사이즈 수 만큼의 데이터를 batch_x, batch_y에 담아 최종 학습할 input, output을 return 해줍니다.
한 epoch이 끝날 때마다 실행하는 함수(on_epoch_end)
사실 이 함수는 없어도 됩니다. 여기서 사용하는 이유는 train에서 데이터의 순서를 shuffle해주기 위함입니다.
self.indices 변수에 전체 데이터의 개수에 맞게 index 배열을 넣어주고, shuffle이 True라면 그 index 배열을 랜덤하게 섞어줍니다. 그럼 위의 __getitem__함수를 호출할 때마다 뒤섞인 데이터가 반환될 수 있습니다.
def on_epoch_end(self):
self.indices = np.arange(len(self.x))
if self.shuffle == True:
np.random.shuffle(self.indices)전체 소스 코드
전체 소스 코드는 다음과 같습니다.
위의 Sequence 로더의 장점은 tensorflow에서 제공하는 fit 함수에서도 사용 할 수 있고, tf.GradientTape 를 사용하는 custom 학습 loop에서도 쉽게 사용할 수 있습니다.
import numpy as np
import math
from tensorflow.keras.utils import Sequence
class Dataloader(Sequence):
def __init__(self, x_set, y_set, batch_size, shuffle=False):
self.x, self.y = x_set, y_set
self.batch_size = batch_size
self.shuffle=shuffle
self.on_epoch_end()
def __len__(self):
return math.ceil(len(self.x) / self.batch_size)
# batch 단위로 직접 묶어줘야 함
def __getitem__(self, idx):
# sampler의 역할(index를 batch_size만큼 sampling해줌)
indices = self.indices[idx*self.batch_size:(idx+1)*self.batch_size]
batch_x = [self.x[i] for i in indices]
batch_y = [self.y[i] for i in indices]
return np.array(batch_x), np.array(batch_y)
# epoch이 끝날때마다 실행
def on_epoch_end(self):
self.indices = np.arange(len(self.x))
if self.shuffle == True:
np.random.shuffle(self.indices)
train_loader = Dataloader(x, y, 128, shuffle=True)
valid_loader = Dataloader(x, y, 128)
test_loader = Dataloader(x, y, 128)
# 방법 1
model.fit(train_loader, validation_data=valid_loader, epochs=10,
workers=4)# multi로 처리할 개수
model.evaluate(test_loader)
# 방법 2
for e in range(epochs):
for x, y in train_loader:
train_step(x, y)
for x, y in valid_loader:
valid_step(x, y)참고자료
http://www.kwangsiklee.com/2018/11/keras%EC%97%90%EC%84%9C-sequence%EB%A5%BC-%EC%9D%B4%EC%9A%A9%ED%95%98%EC%97%AC-%EB%8C%80%EC%9A%A9%EB%9F%89-%EB%8D%B0%EC%9D%B4%ED%84%B0%EC%85%8B-%EC%B2%98%EB%A6%AC%ED%95%98%EA%B8%B0/
https://www.tensorflow.org/api_docs/python/tf/keras/utils/Sequence