Datasets documentation
JAX와 함께 사용
JAX와 함께 사용
이 문서는 datasets를 JAX와 함께 사용하는 방법에 대한 간단한 소개입니다. 특히 데이터셋에서 jax.Array 객체를 어떻게 가져오는지, 그리고 이를 사용하여 JAX 모델을 학습하는 방법에 중점을 둡니다.
위의 코드를 재현하려면
jax와jaxlib가 필요하므로,pip install datasets[jax]로 설치했는지 확인하세요.
데이터셋 형식
기본적으로 datasets는 정수, 실수, 문자열, 리스트 등 일반적인 Python 객체를 반환합니다. JAX는 숫자만 지원하므로 문자열 및 바이너리 객체는 변경되지 않습니다.
numpy와 유사한 JAX 배열을 얻으려면, 데이터셋의 형식을 jax로 설정할 수 있습니다:
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("jax")
>>> ds[0]
{'data': DeviceArray([1, 2], dtype=int32)}
>>> ds[:2]
{'data': DeviceArray([
[1, 2],
[3, 4]], dtype=int32)}Dataset 객체는 Arrow 테이블의 래퍼로, 데이터셋의 배열에서 JAX 배열로의 빠른 읽기를 허용합니다.
DatasetDict 객체에도 동일한 절차가 적용되므로, DatasetDict의 형식을 jax로 설정하면 해당 DatasetDict에 포함된 모든 Dataset이 jax 형식으로 포맷팅됩니다:
>>> from datasets import DatasetDict
>>> data = {"train": {"data": [[1, 2], [3, 4]]}, "test": {"data": [[5, 6], [7, 8]]}}
>>> dds = DatasetDict.from_dict(data)
>>> dds = dds.with_format("jax")
>>> dds["train"][:2]
{'data': DeviceArray([
[1, 2],
[3, 4]], dtype=int32)}또한 고려해야 할 점은 포맷팅이 실제로 데이터에 접근할 때까지 적용되지 않는다는 것입니다. 따라서 데이터셋에서 JAX 배열을 얻으려면 먼저 데이터에 접근해야 하며, 그렇지 않으면 포맷은 그대로 유지됩니다.
마지막으로, 원하는 장치에 데이터를 로드하려면 device 인수를 지정할 수 있습니다. 단, jaxlib.xla_extension.Device는 pickle로도 dill로도 직렬화할 수 없어 지원되지 않으므로, 대신 문자열 식별자를 사용해야 합니다:
>>> import jax
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> device = str(jax.devices()[0]) # Not casting to `str` before passing it to `with_format` will raise a `ValueError`
>>> ds = ds.with_format("jax", device=device)
>>> ds[0]
{'data': DeviceArray([1, 2], dtype=int32)}
>>> ds[0]["data"].device()
TFRT_CPU_0
>>> assert ds[0]["data"].device() == jax.devices()[0]
Truedevice 인수가 with_format에 제공되지 않으면 jax.devices()[0]인 기본 장치가 사용됩니다.
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("jax")
>>> ds[0]
{'data': Array([[1, 2],
[3, 4]], dtype=int32)}>>> 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("jax")
>>> ds[0]
{'data': [Array([1, 2], dtype=int32), Array([3], dtype=int32)]}그러나 이러한 로직은 느린 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("jax")
>>> ds[0]
{'data': Array([[1, 2],
[3, 4]], dtype=int32)}
>>> ds[:2]
{'data': Array([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]], dtype=int32)}기타 기능 유형
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("jax")
>>> ds[:3]
{'label': DeviceArray([0, 0, 1], dtype=int32)}JAX는 숫자만 지원하므로 문자열 및 바이너리 객체는 변경되지 않습니다.
Image 기능 유형을 사용하려면
vision추가 항목을pip install datasets[vision]으로 설치해야 합니다.
>>> from datasets import Dataset, Features, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]["image"].shape
(512, 512, 3)
>>> ds[0]
{'image': DeviceArray([[[ 255, 255, 255],
[ 255, 255, 255],
...,
[ 255, 255, 255],
[ 255, 255, 255]]], dtype=uint8)}
>>> ds[:2]["image"].shape
(2, 512, 512, 3)
>>> ds[:2]
{'image': DeviceArray([[[[ 255, 255, 255],
[ 255, 255, 255],
...,
[ 255, 255, 255],
[ 255, 255, 255]]]], dtype=uint8)}Audio 기능 유형을 사용하려면
audio추가 항목을pip install datasets[audio]로 설치해야 합니다.
>>> from datasets import Dataset, Features, Audio
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]["audio"]["array"]
DeviceArray([-0.059021 , -0.03894043, -0.00735474, ..., 0.0133667 ,
0.01809692, 0.00268555], dtype=float32)
>>> ds[0]["audio"]["sampling_rate"]
DeviceArray(44100, dtype=int32, weak_type=True)데이터 로딩
JAX에는 내장된 데이터 로딩 기능이 없으므로, PyTorch의 DataLoader 또는 TensorFlow의 tf.data.Dataset와 같은 라이브러리를 사용하여 데이터를 로드해야 합니다. 이 주제에 대한 JAX 문서를 인용하면 다음과 같습니다.
“JAX는 프로그램 변환과 가속기 기반 NumPy에 집중하고 있으므로, JAX 라이브러리에 데이터 로딩이나 데이터 전처리를 포함하지 않습니다. 이미 훌륭한 데이터 로더가 많이 있으므로, 새로운 것을 발명하기보다는 기존 것을 사용하는 것이 좋습니다. PyTorch의 데이터 로더를 가져와 NumPy 배열과 함께 작동하도록 작은 shim을 만들겠습니다.”
datasets에서 JAX 포맷팅이 매우 유용한 이유가 바로 이것입니다. HuggingFace Hub의 어떤 모델을 JAX와 함께 사용하더라도 데이터 로딩 부분에 대해 걱정할 필요가 없기 때문입니다.
with_format('jax') 사용
데이터셋에서 JAX 배열을 가져오는 가장 쉬운 방법은 with_format('jax') 메서드를 사용하는 것입니다. HuggingFace Hub의 https://huggingface.co/datasets/ylecun/mnist에 있는 MNIST dataset으로 신경망을 학습하고 있다고 가정해 보겠습니다.
>>> from datasets import load_dataset
>>> ds = load_dataset("ylecun/mnist")
>>> ds = ds.with_format("jax")
>>> ds["train"][0]
{'image': DeviceArray([[ 0, 0, 0, ...],
[ 0, 0, 0, ...],
...,
[ 0, 0, 0, ...],
[ 0, 0, 0, ...]], dtype=uint8),
'label': DeviceArray(5, dtype=int32)}Dataset.iter() 메서드를 사용하여 데이터셋을 배치 단위로 JAX 모델에 전달할 수 있습니다:
>>> for epoch in range(epochs):
... for batch in ds["train"].iter(batch_size=32):
... x, y = batch["image"], batch["label"]
... ...