2차원 음성 대용량 데이터셋을 TfRecord로 만들기

Tensorflow는 pytorch의 dataloader처럼 queue를 사용하여, 전체 데이터셋을 가져온 뒤 그것을 batch 만큼 쪼개서 하는 것이 살짝 번거롭다.

 

즉 이말을 다시 풀어보면, pytorch에서는 dataloader를 사용하여 여러 queue를 사용해서 batch 만큼 데이터셋을 가져온 뒤, 이것을 tensor로 바꿔서 model에 넣는 것이 수월한데 비해

 

Tensorflow에서는 tf.data.Dataset.from_tensor_slices 를 사용해서 전체 데이터셋을 가져오는 예제가 많다.

 

게다가 대용량 데이터셋을 사용하게 될 경우, 데이터셋의 총 사이즈가 3GB 정도가 넘어가면 tensorflow 는 API 관련 오류가 생기면서 data load가 안 될 때가 있다.

 

근 3일 정도 고생하면서 찾아본 정보들을 합쳐서, 음성 데이터셋의 stft 한 결과인 2차원 데이터셋을 tfrecord로 저장하는 방법을 소개한다.

 

 

# 전처리

음성(.wav)파일 모두에 대해, 2차원 stft를 얻었다고 가정하고 진행하겠다.

stft를 바꾸는 방법은 이 블로그 내에 있다.

 

import tensorflow as tf
print(tf.__version__)
import os
import librosa
from glob import glob
import numpy as np


list_inp = sorted(glob('/your/input/dataset/*/*.npz'))
list_tar = sorted(glob('/your/target/dataset/*/*.npz'))

print(len(list_inp))

위의 코드는 모든 음성 파일을 .npz 라는 numpy 형태의 data format으로 미리 저장해둔 상태이고, 이것을 glob으로 가져오는 모습이다. Seq2Seq 모델의 Encoder input인 list_inp와 Decoder input, Real value input인 list_tar의 전체를 가져온다.

 

 

# Writing Tfrecords file

def serialize_example(batch, list1, list2):
    filename = "./train_set.tfrecords"
    writer = tf.io.TFRecordWriter(filename)
    
    for i in range(batch):
        feature = {}
        feature1 = np.load(list1[i])
        feature2 = np.load(list2[i])
        print('feature1 shape {} feature2 shape {}'.format(feature1.shape, feature2.shape)) 
        feature['input'] = tf.train.Feature(float_list=tf.train.FloatList(value=feature1.flatten()))
        feature['target'] = tf.train.Feature(float_list=tf.train.FloatList(value=feature2.flatten()))
        
        features = tf.train.Features(feature=feature)
        example = tf.train.Example(features=features)
        serialized = example.SerializeToString()
        writer.write(serialized)
        print("{}th input {} target {} finished".format(i, list1[i], list2[i]))
        

serialize_example(len(list_inp), list_inp, list_tar)

 

먼저 저장할 tfrecords filename을 정해주고, tensorflow에서 제공해주는 TFRecordWriter로 writer를 선언한다.

for i in range(batch)는 전체 batch 사이즈만큼 반복하는 것이다. 즉 전체 데이터셋을 저장할 것이다. (나의 경우에는 총 8백만개이다.)

 

numpy 형태의 파일을 sorted(glob())으로 가져왔기 때문에, 현재 함수의 인자값인 list1과 list2는 파일 경로이다. 그러므로, np.load를 사용하여 feature1과 feature2에 값을 넣어준다.

 

그 뒤 tfrecords에 저장할 feature['name'] tf.train.Feature(float_list=tf.train.FloatList(value=feature1.flatten()))을 하여 선언해준다.

 

2d numpy 값을 그대로 value=feature1 를 해주면, 오류가 생긴다. 이를 찾아보게 되면, tensorflow에서의 FloatList는 1차원 값만 넣어줄 수 있기 때문이다. 그러므로, flatten()을 하여 넣어준다.

 

나 같은 경우 input과 target을 모두 maximum legnth를 구하고 zeropadding하여 shape을 아는 상태이지만, 만약 데이터셋이 모두 다를 경우 feature['shape'] = tf.train.Feature(int_list=tf.train.IntList(value=feature1.shape) 을 대입하여 나중에 데이터셋을 실제로 model에 넣을 때 shape을 기억하여 변환할 수 있다.

 

그다음 features = tf.train.Features(feature=feature)를 통해 tensorflow의 tensor로 변환해주고,

example 또한 마찬가지로 변환해준다.

 

serialized = example.SerializeToString() 을 통하여 binary? 로 변환해준다. 마지막으로 tf.records 파일을 write 해주는 writer.write(serialized)를 해주면 끝난다.

이것을 batch size만큼 해준다...

 

전체코드는 아래와 같다.

 

 

# 전체 코드

import tensorflow as tf
print(tf.__version__)
import os
import librosa
from glob import glob
import numpy as np


def serialize_example(batch, list1, list2):
    filename = "./train_set.tfrecords"
    writer = tf.io.TFRecordWriter(filename)
    
    for i in range(batch):
        feature = {}
        feature1 = np.load(list1[i])
        feature2 = np.load(list2[i])
        print('feature1 shape {} feature2 shape {}'.format(feature1.shape, feature2.shape)) 
        feature['input'] = tf.train.Feature(float_list=tf.train.FloatList(value=feature1.flatten()))
        feature['target'] = tf.train.Feature(float_list=tf.train.FloatList(value=feature2.flatten()))
        
        features = tf.train.Features(feature=feature)
        example = tf.train.Example(features=features)
        serialized = example.SerializeToString()
        writer.write(serialized)
        print("{}th input {} target {} finished".format(i, list1[i], list2[i]))
        


list_inp = sorted(glob('/your/input/dataset/*/*.npz'))
list_tar = sorted(glob('/your/target/dataset/*/*.npz'))

print(len(list_inp))
serialize_example(len(list_inp), list_inp, list_tar)

 

input, target 각각 8백만개 씩 있는데, 이것들이 과연 tfrecords를 통해 저장하였을 때의 공간적 이득과, 시간이 얼마나 걸리는지 체크해봐야겠다.

 

체크 해본 뒤, Tfrecord 파일을 다시 원래대로 음성 데이터 spectrogram으로 복구하는 것을 이번주내로 올릴 예정이다.