TFRecord 파일 읽고 쓰기

학부 졸업 프로젝트를 수행하며 tensorflow를 통해 딥러닝 모델을 구현하던 중 tensorflow의 학습 데이터 포맷인 tfrecord라는 파일을 만날 수 있었습니다. 딥러닝 모델의 구현을 위해 github을 뒤져보던 중 *.tfrecord라는 새로운 파일 확장자를 만났고, 용량도 별로 되지 않는 것이 모든 데이터를 담고 있다는 것을 보니 머릿속에 “다뤄야 한다”라는 생각이 들었습니다.

그렇다면 tfrecord 파일 포맷은 무엇인가??

tfrecord

tfrecord 파일은 tensorflow의 학습 데이터 등을 저장하기 위한 바이너리 데이터 포맷으로, 구글의 protocol buffer 포맷으로 데이터를 파일에 serialize하여 저장하게 됩니다.

tfrecord 파일의 필요성은 아래와 같습니다. (from. 조대협님의 블로그)
1. csv파일에서와 같이 숫자나 텍스트 데이터를 읽을 때는 크게 지장이 없지만, 이미지 데이터를 읽을 경우 이미지는 jpeg나 png 형태의 파일로 저장되어 있고 이에 대한 메타데이터와 레이블은 별도의 파일에 저장되어 있기 때문에, 학습 데이터를 읽을 때 메타데이터나 레이블 파일 하나만 읽는 것이 아니라 이미지 파일도 별도로 읽어야 하기 때문에 코드가 복잡해진다.
2. 이미지를 jpg난 png 포맷으로 읽어서 매번 디코딩을 하게되면, 그 성능이 저하되서 학습단계에서 데이터를 읽는 부분에서 많은 성능 저하가 발생한다.

이는 저 역시도 딥러닝을 공부하고 모델을 돌려보면서 많이 느꼈던 애로사항입니다. 좋은 하드웨어를 구하지 못해서 구글의 colab 환경을 많이 사용했는데, 데이터를 구글 드라이브에 옮기는 시간은 물론이고 이를 다시 학습 환경으로 가져오는 시간도 길어지는 경우가 종종 있었습니다.

tfrecord 파일 생성

아래의 내용은 tensorflow 공식문서와 ‘peteryuX’님의 얼굴인식 github 레포지토리를 참고했습니다.

tfrecord 파일 생성은 기록하고자 하는 데이터의 feature들을 python dictionary 형태로 정의한 후, 데이터 하나씩 tf.train.Example 객체로 만들어 tf.io.TFRecordWriter를 통해 파일로 저장합니다.

1. 데이터 변환

tensorflow 공식문서에 따르면 tf.train.Example 객체는 아래와 같은 유형의 데이터 타입을 담을 수 있습니다.

tf.train.BytesList
- string
- byte
tf.train.FloatList
- float (float32)
- double (float64)
tf.train.Int64List
- bool
- enum
- int32
- uint32
- int64
- uint64

기록하고자하는 데이터의 타입에 따라 아래와 같은 변환 함수를 사용해 데이터를 기록할 수 있습니다. 예를 들어 문자열 타입의 데이터는 tf.train.BytesList의 객체로 기록하는 것이죠.

  
# The following functions can be used to convert a value to a type compatible
# with tf.train.Example.

def _bytes_feature(value):
    """Returns a bytes_list from a string / byte."""
    if isinstance(value, type(tf.constant(0))):
        value = value.numpy() # BytesList won't unpack a string from an EagerTensor.
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _float_feature(value):
    """Returns a float_list from a float / double."""
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))

def _int64_feature(value):
    """Returns an int64_list from a bool / enum / int / uint."""
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

2. tf.train.Example 객체로 저장

  
def make_example(img_str, source_id, filename):
    # Create a dictionary with features that may be relevant.
    feature = {'image/source_id': _int64_feature(source_id),
               'image/filename': _bytes_feature(filename),
               'image/encoded': _bytes_feature(img_str)}

    return tf.train.Example(features=tf.train.Features(feature=feature))

기록할 데이터들을 python dictionary 형태로 묶어 tf.train.Example 객체로 변환합니다. 위 예제코드에서는 얼굴인식을 위한 데이터를 저장하는 것인데, 각각 source_id는 클래스 ID를, filename은 이미지 파일의 이름, encoded는 이미지 파일 자체를 의미합니다.

3. tf.io.TFRecordWriter로 파일 생성

원본 데이터의 디렉토리 구조는 아래와 같이 일반적인 Classification 데이터의 구조를 가진다고 하겠습니다.

.
|-- 0
|   |-- ****.jpg
|   |-- ****.jpg
|   |-- ...
|
|-- 1
|   |-- ****.jpg
|   |-- ***.jpg
|   |-- ...
|
|-- 2
|   |-- ****.jpg
|   |-- ...
...

아래는 위 디렉토리 구조에서 파일의 이름을 입력받아 tfrecord 파일을 생성하는 코드입니다. 0, 1, 2, ...의 디렉토리 이름을 클래스 ID로 하고 차례로 sample python list에 담아 하나씩 저장합니다. 그 다음 list에 담은 샘플들을 섞고 다시 차례로 *.tfrecord 파일에 적습니다.

  
def main(dataset_path, output_path):
    samples = []
    print("Reading data list...")
    for id_name in tqdm(os.listdir(dataset_path)):
        img_paths = glob(os.path.join(dataset_path, id_name, '*.jpg'))
        for img_path in img_paths:
            filename = os.path.join(id_name, os.path.basename(img_path))
            samples.append((img_path, id_name, filename))
    random.shuffle(samples)

    print("Writing tfrecord file...")
    with tf.io.TFRecordWriter(output_path) as writer:
        for img_path, id_name, filename in tqdm(samples):
            tf_example = make_example(img_str=open(img_path, 'rb').read(),
                                      source_id=int(id_name),
                                      filename=str.encode(filename))
            writer.write(tf_example.SerializeToString())

회귀 문제에 대한 tfrecord 파일 생성 예제를 하나 더 보면 좋을 것 같습니다.

회귀 문제에 대한 label은 .csv나 .txt 같은 별도의 파일에 저장하는 경우가 많습니다. 아래는 이미지 파일명과 그 label 값을 함께 적어놓은 .csv 파일의 예시입니다.

filename,label
0000971160_1, 6.094467065223975
0000971160_2, -8.213548899809783
0000971160_3, -16.811467943863377
0000971160_4, -0.44434946696166655
0000971160_5, -26.681545004531852
0000971160_6, 17.006392102458875
0000971160_7, -18.696981688745513
0000971160_8, 16.358331838072285

위와 같이 저장된 파일도 같은 방식으로 tfrecord 파일로 저장할 수 있습니다. 달라진 점이라면 메타데이터에서 정보를 가져와 활용한다는 것과 그에 맞는 label의 type에 맞게 make_example을 조정했다는 점입니다.

  
def make_example(img_str, label, filename):
    feature = {'encoded': _bytes_feature(img_str),
               'label': _float_feature(label),
               'filename': _bytes_feature(filename)}
    return tf.train.Example(features=tf.train.Features(feature=feature))
    

def main(dataset_path, output_path):
    samples = []
    with open(f"{dataset_path}/metadata.csv") as f:
        f.readline()
        lines = f.readlines()
        for line in tqdm(lines):
            filename, label = line.split(',')
            img_path = os.path.join(dataset_path, mode.split('_')[0], f"{filename}.png")
            label = float(label.replace('\n', ''))

            samples.append((img_path, label, filename))
    random.shuffle(samples)

    with tf.io.TFRecordWriter(output_path) as writer:
        for img_path, label, filename in tqdm(samples):
            tf_example = make_example(img_str=open(img_path, 'rb').read(),
                                      label=label,
                                          filename=str.encode(filename+".png"))
            writer.write(tf_example.SerializeToString())

tfrecord 파일 tf.data.Dataset으로 변환

이렇게 만든 tfrecord 파일을 읽는 방법 역시 간단합니다.

  
def _parse_tfrecord():
    def parse_tfrecord(tfrecord):
        features = {'image/source_id': tf.io.FixedLenFeature([], tf.int64),
                    'image/filename': tf.io.FixedLenFeature([], tf.string),
                    'image/encoded': tf.io.FixedLenFeature([], tf.string)}
        x = tf.io.parse_single_example(tfrecord, features)
        x_train = tf.image.decode_jpeg(x['image/encoded'], channels=3)

        y_train = tf.cast(x['image/source_id'], tf.float32)
        x_train = _transform_images()(x_train)
        y_train = _transform_targets(y_train)
        return (x_train, y_train), y_train
    return parse_tfrecord


def _transform_images():
    def transform_images(x_train):
        x_train = tf.image.resize(x_train, (128, 128))
        x_train = tf.image.random_crop(x_train, (112, 112, 3))
        x_train = tf.image.random_flip_left_right(x_train)
        x_train = tf.image.random_saturation(x_train, 0.6, 1.4)
        x_train = tf.image.random_brightness(x_train, 0.4)
        x_train = x_train / 255
        return x_train
    return transform_images


def _transform_targets(y_train):
    return y_train

기록한 tfrecord 파일을 입력받아 tf.io.parse_single_example()로 하나씩 풀어주면 됩니다. 이렇게 불러온 데이터를 사용하고자하는 모델에 맞게 전처리해주고 알맞은 형태의 return 값을 반환합니다.

  
def load_tfrecord_dataset(tfrecord_name, batch_size, shuffle=True, buffer_size=10240):
    """load dataset from tfrecord"""
    raw_dataset = tf.data.TFRecordDataset(tfrecord_name)
    raw_dataset = raw_dataset.repeat()
    if shuffle:
        raw_dataset = raw_dataset.shuffle(buffer_size=buffer_size)
    dataset = raw_dataset.map(
        _parse_tfrecord(),
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    return dataset

tf.data.TFRecordDataset의 객체를 생성해 .map()을 통해 앞서 정의한 변환과정을 적용해주면 데이터 사용 준비가 마무리됩니다.

tfrecord와 Generator의 속도 비교

그렇다면 python generator를 사용한 tensorflow 데이터셋과 tfrecord를 사용한 tensorflow 데이터셋의 시간 차이를 살펴보겠습니다.

colab 환경에서 진행했고 위에서 소개한 classification 문제에 대한 데이터셋을 비교에 사용했습니다. 총 이미지 수는 9360장, 클래스의 종류는 390가지입니다.

image-description

왼쪽은 tfrecord를 사용한 시간으로 5초 남짓 소요되었지만, 오른쪽 Generator를 통해 디렉토리에서 직접 데이터를 가져오는 방식은 90분 정도의 어마어마한 차이를 보였습니다.

TFRecord 파일 읽고 쓰기

tfrecord

tfrecord 파일 생성

1. 데이터 변환

2. tf.train.Example 객체로 저장

3. tf.io.TFRecordWriter로 파일 생성

tfrecord 파일 tf.data.Dataset으로 변환

tfrecord와 Generator의 속도 비교

Further Reading

[project-review] ArcFace를 활용한 한국인 안면 인식

[paper-review] Scribble-Supervised LiDAR Semantic Segmentation

[paper-review] Denoising diffusion probabilistic models

Trending Tags