TensorFlow와 Datasets 사용

이 문서는 datasets를 TensorFlow와 함께 사용하는 방법에 대한 간단한 소개입니다. 특히 데이터셋에서 tf.Tensor 객체를 가져오는 방법과 Hugging Face Dataset 객체에서 model.fit()과 같은 Keras 메서드로 데이터를 스트리밍하는 방법에 중점을 둡니다.

데이터셋 형식

기본적으로 datasets는 정수, 실수, 문자열, 리스트 등 일반적인 Python 객체를 반환합니다.

TensorFlow 텐서를 얻으려면 데이터셋의 형식을 tf로 설정할 수 있습니다:

>>> from datasets import Dataset
>>> data = [[1, 2],[3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([1, 2])>}
>>> ds[:2]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
array([[1, 2],
       [3, 4]])>}

Dataset 객체는 Arrow 테이블의 래퍼로, 데이터셋의 배열에서 TensorFlow 텐서로의 빠른 읽기를 허용합니다.

데이터셋을 Tensor 객체의 dict로 변환하거나, 이를 통해 TF 샘플을 로드하는 생성기를 작성할 때 유용할 수 있습니다. 데이터셋 전체를 Tensor로 변환하려면 전체 데이터셋을 쿼리하면 됩니다:

>>> ds[:]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
array([[1, 2],
       [3, 4]])>}

N차원 배열

데이터셋이 N차원 배열로 구성되어 있다면, 기본적으로 shape가 고정되어 있으면 동일한 tensor로 간주된다는 것을 확인할 수 있습니다:

>>> from datasets import Dataset
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]  # fixed shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
 array([[1, 2],
        [3, 4]])>}

그렇지 않은 경우, TensorFlow 형식의 데이터셋은 단일 텐서 대신 RaggedTensor를 출력합니다:

>>> from datasets import Dataset
>>> data = [[[1, 2],[3]],[[4, 5, 6],[7, 8]]]  # varying shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': <tf.RaggedTensor [[1, 2], [3]]>}

그러나 이러한 로직은 느린 shape 비교와 데이터 복사를 요구하는 경우가 많습니다. 이를 피하려면 Array feature type을 명시적으로 사용하고 텐서의 shape를 지정해야 합니다:

>>> from datasets import Dataset, Features, Array2D
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]
>>> features = Features({"data": Array2D(shape=(2, 2), dtype='int32')})
>>> ds = Dataset.from_dict({"data": data}, features=features)
>>> ds = ds.with_format("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
 array([[1, 2],
        [3, 4]])>}
>>> ds[:2]
{'data': <tf.Tensor: shape=(2, 2, 2), dtype=int64, numpy=
 array([[[1, 2],
         [3, 4]],
 
        [[5, 6],
         [7, 8]]])>}

기타 기능 유형

ClassLabel 데이터가 텐서로 올바르게 변환됩니다:

>>> from datasets import Dataset, Features, ClassLabel
>>> labels = [0, 0, 1]
>>> features = Features({"label": ClassLabel(names=["negative", "positive"])})
>>> ds = Dataset.from_dict({"label": labels}, features=features) 
>>> ds = ds.with_format("tf")  
>>> ds[:3]
{'label': <tf.Tensor: shape=(3,), dtype=int64, numpy=array([0, 0, 1])>}

문자열과 바이너리 객체도 지원됩니다:

>>> from datasets import Dataset, Features 
>>> text = ["foo", "bar"]
>>> data = [0, 1] 
>>> ds = Dataset.from_dict({"text": text, "data": data})  
>>> ds = ds.with_format("tf") 
>>> ds[:2]
{'text': <tf.Tensor: shape=(2,), dtype=string, numpy=array([b'foo', b'bar'], dtype=object)>,
 'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([0, 1])>}

특정 열만 명시적으로 포맷하고 나머지 열은 포맷하지 않은 상태로 유지할 수도 있습니다:

>>> ds = ds.with_format("tf", columns=["data"], output_all_columns=True)
>>> ds[:2]
{'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([0, 1])>,
 'text': ['foo', 'bar']}

PyTorch는 숫자만 지원하므로 문자열 및 바이너리 객체는 변경되지 않습니다.

Image와 Audio 기능 유형도 지원됩니다.

Image 기능 유형을 사용하려면 vision 추가 항목을 pip install datasets[vision]으로 설치해야 합니다.

>>> from datasets import Dataset, Features, Audio, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features) 
>>> ds = ds.with_format("tf")  
>>> ds[0]
{'image': <tf.Tensor: shape=(512, 512, 4), dtype=uint8, numpy=
 array([[[255, 215, 106, 255],
         [255, 215, 106, 255],
         ...,
         [255, 255, 255, 255],
         [255, 255, 255, 255]]], dtype=uint8)>}
>>> ds[:2]
{'image': <tf.Tensor: shape=(2, 512, 512, 4), dtype=uint8, numpy=
 array([[[[255, 215, 106, 255],
          [255, 215, 106, 255],
          ...,
          [255, 255, 255, 255],
          [255, 255, 255, 255]]]], dtype=uint8)>}

Audio 기능 유형을 사용하려면 audio 추가 항목을 pip install datasets[audio]로 설치해야 합니다.

>>> from datasets import Dataset, Features, Audio, Image
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features) 
>>> ds = ds.with_format("tf")  
>>> ds[0]["audio"]["array"]
<tf.Tensor: shape=(202311,), dtype=float32, numpy=
array([ 6.1035156e-05,  1.5258789e-05,  1.6784668e-04, ...,
       -1.5258789e-05, -1.5258789e-05,  1.5258789e-05], dtype=float32)>
>>> ds[0]["audio"]["sampling_rate"]
<tf.Tensor: shape=(), dtype=int32, numpy=44100>

데이터 로딩

데이터셋에 인덱싱하여 개별 샘플과 배치만 로드하는 것은 가능하지만, fit()와 predict()와 같은 Keras 메서드를 사용하려는 경우 이 방법은 작동하지 않습니다. 데이터셋에서 배치를 셔플하고 로드하는 생성자 함수를 작성하여 fit()을 실행할 수도 있지만, 이는 불필요한 작업이 많아 보입니다. 대신 데이터셋에서 데이터를 온더플라이로 스트리밍하려면 to_tf_dataset() 메서드를 사용하여 데이터셋을 tf.data.Dataset으로 변환하는 것을 권장합니다.

tf.data.Dataset 클래스는 다양한 사용 사례를 지원합니다. 메모리의 Tensor에서 생성되거나, 디스크 또는 외부 저장소의 파일을 읽기 위한 로드 함수를 사용하여 생성되는 경우가 많습니다. 데이터셋은 map() 메서드로 임의로 변환할 수 있으며, batch() 및 shuffle()와 같은 메서드를 사용하여 학습에 바로 사용할 수 있는 데이터셋을 만들 수 있습니다. 이러한 메서드는 저장된 데이터를 어떤 방식으로든 수정하지 않습니다. 대신, 데이터셋이 반복될 때(보통 모델 학습 또는 추론 중에) 실행될 데이터 파이프라인 그래프를 구성합니다. 이는 Hugging Face Dataset 객체의 map() 메서드와 다르며, 후자는 map 함수를 즉시 실행하여 새로 생성되거나 변경된 컬럼을 저장합니다.

전체 데이터 전처리 파이프라인을 tf.data.Dataset으로 컴파일할 수 있으므로, 이 접근 방식은 대규모 병렬 및 비동기 데이터 로딩과 학습을 가능하게 합니다. 그러나 그래프 컴파일 요구 사항은 제한 사항이 될 수 있으며, 특히 TF 그래프의 일부로 (아직) 컴파일할 수 없는 경우가 대부분인 Hugging Face 토크나이저에서 그렇습니다. 그 결과, 임의의 Python 함수를 사용할 수 있는 Hugging Face dataset으로 데이터를 사전 처리한 후, to_tf_dataset()를 사용하여 tf.data.Dataset으로 변환하여 학습에 바로 사용할 수 있는 배치된 dataset을 얻는 것을 일반적으로 권장합니다. 이 접근 방식의 예제는 transformers의 예제 또는 노트북을 참고하세요.

to_tf_dataset() 사용

to_tf_dataset() 사용은 간단합니다. 데이터셋이 전처리되어 준비되면 다음과 같이 호출하면 됩니다:

>>> from datasets import Dataset
>>> data = {"inputs": [[1, 2],[3, 4]], "labels": [0, 1]}
>>> ds = Dataset.from_dict(data)
>>> tf_ds = ds.to_tf_dataset(
            columns=["inputs"],
            label_cols=["labels"],
            batch_size=2,
            shuffle=True
            )

여기 반환된 tf_ds 객체는 이제 학습에 완전히 준비된 상태이며, model.fit()에 직접 전달할 수 있습니다. 데이터셋을 생성할 때 배치 크기를 설정했으므로 fit()을 호출할 때 이를 지정할 필요가 없습니다:

>>> model.fit(tf_ds, epochs=2)

인수에 대한 전체 설명은 to_tf_dataset() 문서를 참고하세요. 많은 경우 호출에 collate_fn을 추가해야 합니다. 이는 데이터셋의 여러 요소를 받아 하나의 배치로 결합하는 함수입니다. 모든 요소의 길이가 동일한 경우 내장 기본 collator만으로도 충분하지만, 더 복잡한 작업에는 사용자 정의 collator가 필요할 수 있습니다. 특히, 다양한 시퀀스 길이를 가진 샘플이 있는 많은 작업에서는 배치를 올바르게 패딩할 수 있는 data collator가 필요합니다. 가변 시퀀스 길이가 매우 흔한 transformers NLP examples와 notebooks에서 이에 대한 예시를 확인할 수 있습니다.

to_tf_dataset로 로딩하는 것이 느리다면, num_workers 인수를 사용할 수도 있습니다. 이 기능은 여러 서브프로세스를 시작하여 데이터를 병렬로 로딩합니다. 이 기능은 비교적 최근에 추가되어 아직 실험적인 단계이므로, 사용 중 버그를 발견하시면 이슈를 등록해 주세요!

to_tf_dataset을 사용할 때

독자라면 이 시점에서 동일한 목표를 달성하기 위한 두 가지 접근법이 제시되었음을 눈치챘을 것입니다. TensorFlow 모델에 데이터셋을 전달하려면 .with_format('tf')를 사용하여 데이터셋을 Tensor 또는 Tensors의 dict로 변환하거나, to_tf_dataset()를 사용하여 데이터셋을 tf.data.Dataset로 변환할 수 있습니다. 이 두 가지 모두 model.fit()에 전달할 수 있으므로, 어떤 것을 선택해야 할까요?

핵심은 전체 데이터셋을 Tensor로 변환하면 정적(static)이며 RAM에 완전히 로드된다는 점입니다. 이는 간단하고 편리하지만, 다음 중 하나라도 해당된다면 to_tf_dataset()를 사용하는 것이 좋습니다:

• 데이터셋이 RAM에 담기에는 너무 큽니다. to_tf_dataset()는 한 번에 배치 하나만 스트리밍하므로, 이 방법으로 매우 큰 데이터셋도 처리할 수 있습니다.
• dataset.with_transform() 또는 collate_fn를 사용하여 무작위 변환을 적용하려는 경우입니다. 이는 비전 모델 학습 시 이미지 증강이나 masked language model 학습 시 무작위 마스킹과 같이 여러 모달리티에서 흔히 사용됩니다. to_tf_dataset()를 사용하면 배치 로딩 시점에 해당 변환이 적용되므로, 동일한 샘플이 로드될 때마다 서로 다른 증강이 적용됩니다. 이는 일반적으로 원하는 동작입니다.
• 데이터의 차원이 가변적일 수 있습니다. 예를 들어 NLP의 입력 텍스트는 토큰 수가 서로 다릅니다. 가변 차원을 가진 샘플로 배치 생성 시 표준적인 해결책은 짧은 샘플을 가장 긴 샘플의 길이에 맞게 패딩하는 것입니다. to_tf_dataset를 통해 데이터셋에서 샘플을 스트리밍할 때, collate_fn을 통해 각 배치에 이 패딩을 적용할 수 있습니다. 그러나 이러한 데이터셋을 밀집 Tensor로 변환하려면 데이터셋 전체에서 가장 긴 샘플의 길이에 맞게 샘플을 패딩해야 합니다. 이로 인해 막대한 양의 패딩이 발생하여 메모리가 낭비되고 모델의 속도가 저하될 수 있습니다.

주의 사항 및 제한 사항

현재 to_tf_dataset()는 항상 배치된 데이터셋을 반환합니다 - 곧 배치되지 않은 데이터셋에 대한 지원을 추가할 예정입니다!