Sử dụng với PyTorch

Tài liệu này là phần giới thiệu nhanh về cách sử dụng datasets với PyTorch, tập trung đặc biệt vào cách lấy các đối tượng torch.Tensor từ các dataset của chúng ta, và cách sử dụng DataLoader của PyTorch cùng với Dataset của Hugging Face để đạt hiệu suất tốt nhất.

Đị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.

Để nhận các tensor PyTorch thay thế, bạn có thể đặt định dạng của dataset thành pytorch bằng cách sử dụng Dataset.with_format():

>>> from datasets import Dataset
>>> data = [[1, 2],[3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([1, 2])}
>>> ds[:2]
{'data': tensor([[1, 2],
         [3, 4]])}

Một đối tượng Dataset là lớp bọc của một bảng Arrow, cho phép đọc nhanh không sao chép (zero-copy) từ các mảng trong bộ dữ liệu sang các tensor PyTorch.

Để tải dữ liệu dưới dạng tensor trên GPU, hãy chỉ định đối số device:

>>> import torch
>>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
>>> ds = ds.with_format("torch", device=device)
>>> ds[0]
{'data': tensor([1, 2], device='cuda: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("torch")
>>> ds[0]
{'data': tensor([[1, 2],
         [3, 4]])}

>>> 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': [tensor([1, 2]), tensor([3])]}

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("torch")
>>> ds[0]
{'data': tensor([[1, 2],
         [3, 4]])}
>>> ds[:2]
{'data': tensor([[[1, 2],
          [3, 4]],
 
         [[5, 6],
          [7, 8]]])}

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

Dữ liệu ClassLabel được chuyển đổi thành tensor 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("torch")  
>>> ds[:3]
{'label': tensor([0, 0, 1])}

Các đối tượng chuỗi và nhị phân không thay đổi, vì PyTorch 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, Audio, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features) 
>>> ds = ds.with_format("torch")
>>> ds[0]["image"].shape
torch.Size([512, 512, 4])
>>> ds[0]
{'image': tensor([[[255, 215, 106, 255],
         [255, 215, 106, 255],
         ...,
         [255, 255, 255, 255],
         [255, 255, 255, 255]]], dtype=torch.uint8)}
>>> ds[:2]["image"].shape
torch.Size([2, 512, 512, 4])
>>> ds[:2]
{'image': tensor([[[[255, 215, 106, 255],
          [255, 215, 106, 255],
          ...,
          [255, 255, 255, 255],
          [255, 255, 255, 255]]]], dtype=torch.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, Image
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features) 
>>> ds = ds.with_format("torch")  
>>> ds[0]["audio"]["array"]
tensor([ 6.1035e-05,  1.5259e-05,  1.6785e-04,  ..., -1.5259e-05,
        -1.5259e-05,  1.5259e-05])
>>> ds[0]["audio"]["sampling_rate"]
tensor(44100)

Tải dữ liệu

Tương tự như các đối tượng torch.utils.data.Dataset, một Dataset có thể được truyền trực tiếp vào một DataLoader của PyTorch:

>>> import numpy as np
>>> from datasets import Dataset 
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(16)
>>> label = np.random.randint(0, 2, size=16)
>>> ds = Dataset.from_dict({"data": data, "label": label}).with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=4)
>>> for batch in dataloader:
...     print(batch)                                                                                            
{'data': tensor([0.0047, 0.4979, 0.6726, 0.8105]), 'label': tensor([0, 1, 0, 1])}
{'data': tensor([0.4832, 0.2723, 0.4259, 0.2224]), 'label': tensor([0, 0, 0, 0])}
{'data': tensor([0.5837, 0.3444, 0.4658, 0.6417]), 'label': tensor([0, 1, 0, 0])}
{'data': tensor([0.7022, 0.1225, 0.7228, 0.8259]), 'label': tensor([1, 1, 1, 1])}

Tối ưu hóa việc tải dữ liệu

Có một số cách bạn có thể tăng tốc độ tải dữ liệu, giúp tiết kiệm thời gian, đặc biệt khi làm việc với các bộ dữ liệu lớn. PyTorch cung cấp việc tải dữ liệu song song hóa, lấy các lô chỉ mục thay vì từng cái một, và streaming để lặp qua bộ dữ liệu mà không cần tải xuống đĩa.

Sử dụng nhiều Worker

Bạn có thể song song hóa việc tải dữ liệu bằng đối số num_workers của DataLoader PyTorch và đạt được thông lượng cao hơn.

Bên trong, DataLoader khởi động num_workers tiến trình. Mỗi tiến trình nạp lại bộ dữ liệu được truyền vào DataLoader và được dùng để truy vấn các ví dụ. Việc nạp lại bộ dữ liệu bên trong một worker không làm đầy RAM của bạn, vì nó chỉ ánh xạ bộ nhớ bộ dữ liệu lại từ ổ đĩa của bạn.

>>> import numpy as np
>>> from datasets import Dataset, load_from_disk
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).save_to_disk("my_dataset")
>>> ds = load_from_disk("my_dataset").with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=32, num_workers=4)

Dòng dữ liệu

Dòng chảy một dataset bằng cách tải nó dưới dạng IterableDataset. Điều này cho phép bạn lặp qua một dataset từ xa một cách tuần tự mà không cần tải nó xuống đĩa, hoặc lặp qua các tệp dữ liệu cục bộ. Tìm hiểu thêm về loại dataset nào phù hợp nhất cho trường hợp sử dụng của bạn trong hướng dẫn chọn giữa dataset thông thường và dataset dạng iterable.

Một tập dữ liệu có thể lặp lại từ datasets kế thừa từ torch.utils.data.IterableDataset nên bạn có thể truyền nó vào một torch.utils.data.DataLoader:

>>> import numpy as np
>>> from datasets import Dataset, load_dataset
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).push_to_hub("<username>/my_dataset")  # Upload to the Hugging Face Hub
>>> my_iterable_dataset = load_dataset("<username>/my_dataset", streaming=True, split="train")
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32)

Nếu tập dữ liệu được chia thành nhiều phần (tức là tập dữ liệu bao gồm nhiều tệp dữ liệu), thì bạn có thể stream song song bằng cách sử dụng num_workers:

>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> my_iterable_dataset.num_shards
39
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32, num_workers=4)

Trong trường hợp này, mỗi worker được giao một tập con của danh sách các shard để stream từ đó.

Kiểm tra điểm và tiếp tục

Nếu bạn cần một DataLoader có thể checkpoint và tiếp tục ở giữa quá trình huấn luyện, bạn có thể sử dụng StatefulDataLoader từ torchdata:

>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(my_iterable_dataset, batch_size=32, num_workers=4)
>>> # save in the middle of training
>>> state_dict = dataloader.state_dict()
>>> # and resume later
>>> dataloader.load_state_dict(state_dict)

Điều này có thể thực hiện được nhờ IterableDataset.state_dict() và IterableDataset.load_state_dict().

Phân tán

Để chia bộ dữ liệu của bạn trên các nút huấn luyện, bạn có thể sử dụng datasets.distributed.split_dataset_by_node():

import os
from datasets.distributed import split_dataset_by_node

ds = split_dataset_by_node(ds, rank=int(os.environ["RANK"]), world_size=int(os.environ["WORLD_SIZE"]))

Điều này hoạt động cho cả dataset kiểu map và dataset kiểu iterable. Dataset được chia cho node có rank rank trong một nhóm các node có kích thước world_size.

Đối với các bộ dữ liệu theo phong cách map:

Mỗi node được gán một phần dữ liệu, ví dụ rank 0 nhận phần đầu tiên của tập dữ liệu.

Đối với các tập dữ liệu có thể lặp lại:

Chiến lược strategy phân tách có thể là "auto", "shards" hoặc "examples":

• "shards": các shard được phân bổ đều giữa các node
• "examples": mỗi nút giữ lại 1 ví dụ trong số world_size, bỏ qua các ví dụ còn lại
• "auto" (mặc định): sử dụng "shards" nếu số lượng shards của dataset là một ước số của world_size (tức là nếu dataset.num_shards % world_size == 0), đây là phương án được tối ưu nhất. Ngược lại, sử dụng "examples".

Điều này cũng có thể được kết hợp với torch.utils.data.DataLoader nếu bạn muốn mỗi node sử dụng nhiều worker để tải dữ liệu.

Nếu bạn xáo trộn bộ dữ liệu dạng lặp (iterable dataset) trong một cấu hình phân tán, hãy đảm bảo đặt một seed cố định trong IterableDataset.shuffle() để danh sách các shard đã xáo trộn giống nhau được sử dụng trên mọi nút, từ đó biết được nút nào nên bỏ qua những shard nào.

Song song dữ liệu phân mảnh lai (HSDP)

Một torch.distributed.device_mesh.DeviceMesh mô tả topology huấn luyện của bạn dưới dạng một lưới các rank, mỗi chiều tương ứng với một chiến lược song song. Chỉ các chiều song song dữ liệu mới quan trọng đối với việc nạp dữ liệu, và việc chia sẻ hỗn hợp khiến chúng dễ bị bỏ sót:

• "fsdp" (thường là trong cùng một node): nhóm FSDP mà ở đó mỗi rank chỉ giữ một shard của các trạng thái mô hình (tham số, gradient và trạng thái optimizer) và tạm thời all-gather toàn bộ tham số của một module ngay trước khi tính toán trên module đó. Quá trình huấn luyện vẫn là data-parallel: mỗi rank chạy một lần forward và backward hoàn chỉnh trên batch riêng của nó, do đó chiều này cần dữ liệu khác nhau cho từng rank, giống như DDP.
• "dp_replicate" (thường là giữa các node): các bản sao trùng lặp của mô hình đã được chia mảnh đó, được đồng bộ hóa bằng một thao tác all-reduce bổ sung giữa các replica. Các replica cũng xử lý các dữ liệu khác nhau.
• "tp" (intra-node): các rank tensor parallel chia nhỏ trọng số của mỗi lớp để cùng tính toán một batch. Chúng không phải là data-parallel và phải đọc cùng các ví dụ.

Độ song song dữ liệu là tích của các chiều phân mảnh (FSDP) và nhân bản, nhỏ hơn WORLD_SIZE ngay khi lưới có một chiều như "tp". Truyền tích đó và rank được làm phẳng trên hai chiều đó cho datasets.distributed.split_dataset_by_node():

from torch.distributed.device_mesh import init_device_mesh
from datasets.distributed import split_dataset_by_node

# Hybrid sharding spreads data parallelism over two mesh dimensions, but both
# need different data per rank, so `dp_size` is their product.
mesh = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("dp_replicate", "fsdp", "tp"))
dp_replicate_mesh, fsdp_mesh = mesh["dp_replicate"], mesh["fsdp"]

dp_size = dp_replicate_mesh.size() * fsdp_mesh.size()
dp_rank = dp_replicate_mesh.get_local_rank() * shard.size() + fsdp_mesh.get_local_rank()

ds = split_dataset_by_node(ds, rank=dp_rank, world_size=dp_size)

Ở đây dp_size là 2 * 2 = 4: mỗi rank song song dữ liệu đọc một phần tư của bộ dữ liệu, và hai rank của mỗi cặp "tp" đọc cùng các ví dụ.