Datasets

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v4.8.4).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Sử dụng với JAX

Tài liệu này là phần giới thiệu nhanh về cách sử dụng datasets với JAX, tập trung đặc biệt vào cách lấy các đối tượng jax.Array từ các dataset của chúng ta và cách sử dụng chúng để huấn luyện các mô hình JAX.

jax và jaxlib là bắt buộc để tái tạo mã ở trên, vì vậy hãy đảm bảo bạn cài đặt chúng bằng pip install datasets[jax].

Định dạng tập dữ liệu

Mặc định, các tập dữ liệu trả về các đối tượng Python thông thường: số nguyên, số thực, chuỗi, danh sách, v.v., và các đối tượng chuỗi và nhị phân không thay đổi, vì JAX chỉ hỗ trợ số.

Để nhận các mảng JAX (tương tự numpy) thay vào đó, bạn có thể đặt định dạng của tập dữ liệu thành 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)}

Một đối tượng Dataset là lớp bọc của một bảng Arrow, cho phép đọc nhanh từ các mảng trong bộ dữ liệu sang các mảng JAX.

Lưu ý rằng quy trình hoàn toàn tương tự cũng áp dụng cho các đối tượng DatasetDict, do đó khi đặt định dạng của một DatasetDict thành jax, tất cả các Dataset bên trong sẽ được định dạng thành 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)}

Một điều khác bạn cần lưu ý là định dạng chỉ được áp dụng khi bạn thực sự truy cập dữ liệu. Vì vậy, nếu bạn muốn lấy một mảng JAX từ một dataset, bạn cần truy cập dữ liệu trước, nếu không định dạng sẽ không thay đổi.

Cuối cùng, để nạp dữ liệu vào thiết bị bạn mong muốn, bạn có thể chỉ định đối số device, nhưng lưu ý rằng jaxlib.xla_extension.Device không được hỗ trợ vì nó không thể được serialize bằng cả pickle lẫn dill, do đó bạn sẽ cần sử dụng định danh dạng chuỗi của nó thay thế:

>>> 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]
True

Lưu ý rằng nếu không cung cấp đối số device cho with_format thì nó sẽ sử dụng thiết bị mặc định là jax.devices()[0].

Mảng N chiều

Nếu tập dữ liệu của bạn bao gồm các mảng N chiều, bạn sẽ thấy rằng mặc định chúng được coi là cùng một tensor nếu hình dạng được cố định:

>>> 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)]}

Tuy nhiên, logic này thường đòi hỏi các phép so sánh shape và sao chép dữ liệu chậm. Để tránh điều này, bạn phải sử dụng rõ ràng kiểu đặc trưng Array và chỉ định shape của các tensor của bạn:

>>> 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)}

Các loại tính năng khác

Dữ liệu ClassLabel được chuyển đổi thành mảng một cách chính xác:

>>> 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)}

Các đối tượng chuỗi và nhị phân không thay đổi, vì JAX chỉ hỗ trợ các số.

Các loại tính năng Image và Audio cũng được hỗ trợ.

Để sử dụng kiểu tính năng Image, bạn cần cài đặt phần mở rộng vision bằng lệnh 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)}

Để sử dụng kiểu tính năng Audio, bạn cần cài đặt phần mở rộng audio bằng lệnh 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)

Tải dữ liệu

JAX không có khả năng tải dữ liệu tích hợp sẵn, vì vậy bạn sẽ cần sử dụng một thư viện như PyTorch để tải dữ liệu của bạn bằng DataLoader hoặc TensorFlow bằng tf.data.Dataset. Trích dẫn từ tài liệu JAX về chủ đề này: “JAX tập trung hoàn toàn vào việc chuyển đổi chương trình và NumPy có hỗ trợ từ bộ tăng tốc, vì vậy chúng tôi không bao gồm việc tải hoặc xử lý dữ liệu trong thư viện JAX. Đã có rất nhiều trình tải dữ liệu tuyệt vời ngoài kia, vì vậy hãy chỉ sử dụng chúng thay vì phát minh lại bất cứ thứ gì. Chúng tôi sẽ lấy trình tải dữ liệu của PyTorch và tạo một shim nhỏ để nó hoạt động với các mảng NumPy.”.

Đó là lý do tại sao việc định dạng JAX trong datasets lại hữu ích đến vậy, vì nó cho phép bạn sử dụng bất kỳ mô hình nào từ HuggingFace Hub với JAX, mà không cần phải lo lắng về phần tải dữ liệu.

Sử dụng với with_format('jax')

Cách dễ nhất để lấy các mảng JAX từ một dataset là sử dụng phương thức with_format('jax'). Giả sử chúng ta muốn huấn luyện một mạng nơ-ron trên dataset MNIST có sẵn trên HuggingFace Hub tại https://huggingface.co/datasets/ylecun/mnist.

>>> 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)}

Khi định dạng đã được thiết lập, chúng ta có thể cung cấp bộ dữ liệu cho mô hình JAX theo từng lô bằng cách sử dụng phương thức Dataset.iter():

>>> for epoch in range(epochs):
...     for batch in ds["train"].iter(batch_size=32):
...         x, y = batch["image"], batch["label"]
...         ...

Update on GitHub

←Sử dụng với NumPy Sử dụng với Pandas→