Sự khác nhau giữa Dataset và IterableDataset

Có hai loại đối tượng dataset, một Dataset và một IterableDataset. Việc bạn chọn sử dụng hay tạo loại dataset nào phụ thuộc vào kích thước của dataset. Nhìn chung, IterableDataset là lý tưởng cho các dataset lớn (hàng trăm GB!) nhờ hành vi lazy và lợi thế về tốc độ, trong khi Dataset rất phù hợp cho mọi trường hợp còn lại. Trang này sẽ so sánh sự khác biệt giữa Dataset và IterableDataset để giúp bạn chọn đối tượng dataset phù hợp.

Tải xuống và phát trực tiếp

Khi bạn có một Dataset thông thường, bạn có thể truy cập nó bằng my_dataset[0]. Điều này cho phép truy cập ngẫu nhiên vào các hàng. Các dataset như vậy còn được gọi là dataset kiểu “map”. Ví dụ, bạn có thể tải xuống ImageNet-1k như sau và truy cập bất kỳ hàng nào:

from datasets import load_dataset

imagenet = load_dataset("timm/imagenet-1k-wds", split="train")  # downloads the full dataset
print(imagenet[0])

Tuy nhiên, một lưu ý là bạn phải lưu toàn bộ bộ dữ liệu trên ổ đĩa hoặc trong bộ nhớ, điều này ngăn bạn truy cập các bộ dữ liệu lớn hơn dung lượng ổ đĩa. Vì điều này có thể gây bất tiện với các bộ dữ liệu lớn, nên có một loại bộ dữ liệu khác, đó là IterableDataset. Khi bạn có một IterableDataset, bạn có thể truy cập nó bằng một vòng lặp for để tải dữ liệu theo từng phần khi bạn lặp qua bộ dữ liệu. Nhờ đó, chỉ một phần nhỏ các ví dụ được tải vào bộ nhớ, và bạn không ghi bất cứ thứ gì lên ổ đĩa.

Ví dụ, bạn có thể stream bộ dữ liệu ImageNet-1k mà không cần tải nó xuống đĩa:

from datasets import load_dataset

imagenet = load_dataset("timm/imagenet-1k-wds", split="train", streaming=True)  # will start loading the data when iterated over
for example in imagenet:
    print(example)
    break

Streaming có thể đọc dữ liệu trực tuyến mà không cần ghi bất kỳ tệp nào xuống đĩa. Ví dụ, bạn có thể stream các dataset được tạo từ nhiều shard, mỗi shard có dung lượng hàng trăm gigabyte như C4 hoặc LAION-2B. Tìm hiểu thêm về cách stream một dataset trong Dataset Streaming Guide.

Tuy nhiên, đây không phải là điểm khác biệt duy nhất, vì hành vi “lazy” của IterableDataset cũng xuất hiện trong quá trình tạo và xử lý dataset.

Tạo các bộ dữ liệu dạng map và các bộ dữ liệu dạng lặp

Bạn có thể tạo một Dataset bằng cách sử dụng danh sách hoặc từ điển, và dữ liệu sẽ được chuyển đổi hoàn toàn sang Arrow để bạn có thể dễ dàng truy cập bất kỳ hàng nào:

my_dataset = Dataset.from_dict({"col_1": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]})
print(my_dataset[0])

Mặt khác, để tạo một IterableDataset, bạn phải cung cấp một cách “lười” để tải dữ liệu. Trong Python, chúng ta thường sử dụng các hàm generator. Các hàm này yield từng ví dụ một, điều này có nghĩa là bạn không thể truy cập một hàng bằng cách cắt nó như một Dataset thông thường:

def my_generator(n):
    for i in range(n):
        yield {"col_1": i}

my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs={"n": 10})
for example in my_iterable_dataset:
    print(example)
    break

Tải các tệp cục bộ hoàn toàn và theo từng phần

Bạn có thể chuyển đổi các tệp dữ liệu cục bộ hoặc từ xa thành Dataset Arrow bằng cách sử dụng load_dataset():

data_files = {"train": ["path/to/data.csv"]}
my_dataset = load_dataset("csv", data_files=data_files, split="train")
print(my_dataset[0])

Tuy nhiên, điều này đòi hỏi một bước chuyển đổi từ định dạng CSV sang Arrow, điều này sẽ tốn thời gian và dung lượng ổ đĩa nếu bộ dữ liệu của bạn lớn.

Để tiết kiệm dung lượng ổ đĩa và bỏ qua bước chuyển đổi, bạn có thể định nghĩa một IterableDataset bằng cách stream trực tiếp từ các tệp cục bộ. Nhờ đó, dữ liệu sẽ được đọc dần từ các tệp cục bộ khi bạn lặp qua tập dữ liệu:

data_files = {"train": ["path/to/data.csv"]}
my_iterable_dataset = load_dataset("csv", data_files=data_files, split="train", streaming=True)
for example in my_iterable_dataset:  # this reads the CSV file progressively as you iterate over the dataset
    print(example)
    break

Nhiều định dạng tệp được hỗ trợ, chẳng hạn như CSV, JSONL và Parquet, cũng như tệp hình ảnh và âm thanh. Bạn có thể tìm thêm thông tin trong các hướng dẫn tương ứng để tải các bộ dữ liệu bảng, văn bản, thị giác và âm thanh.

Xử lý dữ liệu theo kiểu Eager và xử lý dữ liệu theo kiểu Lazy

Khi bạn xử lý một đối tượng Dataset bằng Dataset.map(), toàn bộ tập dữ liệu sẽ được xử lý ngay lập tức và trả về. Điều này tương tự như cách pandas hoạt động, chẳng hạn.

my_dataset = my_dataset.map(process_fn)  # process_fn is applied on all the examples of the dataset
print(my_dataset[0])

Mặt khác, do bản chất “lười” của IterableDataset, việc gọi IterableDataset.map() sẽ không áp dụng hàm map của bạn lên toàn bộ tập dữ liệu. Thay vào đó, hàm map của bạn được áp dụng theo thời gian thực.

Vì vậy, bạn có thể nối nhiều bước xử lý lại với nhau và tất cả chúng sẽ được thực hiện đồng thời khi bạn bắt đầu lặp qua tập dữ liệu:

my_iterable_dataset = my_iterable_dataset.map(process_fn_1)
my_iterable_dataset = my_iterable_dataset.filter(filter_fn)
my_iterable_dataset = my_iterable_dataset.map(process_fn_2)

# process_fn_1, filter_fn and process_fn_2 are applied on-the-fly when iterating over the dataset
for example in my_iterable_dataset:  
    print(example)
    break

Xáo trộn chính xác và xáo trộn gần đúng nhanh

Khi bạn trộn một Dataset bằng Dataset.shuffle(), bạn sẽ áp dụng một phép trộn chính xác cho bộ dữ liệu. Nó hoạt động bằng cách lấy một danh sách các chỉ số [0, 1, 2, ... len(my_dataset) - 1] và trộn danh sách này. Sau đó, truy cập my_dataset[0] sẽ trả về hàng và chỉ số được xác định bởi phần tử đầu tiên của ánh xạ chỉ số đã được trộn:

my_dataset = my_dataset.shuffle(seed=42)
print(my_dataset[0])

Vì chúng ta không có khả năng truy cập ngẫu nhiên vào các hàng trong trường hợp của IterableDataset, nên chúng ta không thể sử dụng một danh sách các chỉ số đã xáo trộn và truy cập một hàng ở một vị trí tùy ý. Điều này ngăn cản việc sử dụng xáo trộn chính xác. Thay vào đó, một phương pháp xáo trộn gần đúng nhanh được sử dụng trong IterableDataset.shuffle(). Nó sử dụng một bộ đệm xáo trộn để lấy mẫu các ví dụ ngẫu nhiên một cách lặp đi lặp lại từ tập dữ liệu. Vì tập dữ liệu vẫn được đọc theo cách lặp, nó cung cấp hiệu suất tốc độ tuyệt vời:

my_iterable_dataset = my_iterable_dataset.shuffle(seed=42, buffer_size=100)
for example in my_iterable_dataset:
    print(example)
    break

Tuy nhiên, việc sử dụng bộ đệm trộn (shuffle buffer) là chưa đủ để cung cấp một quá trình trộn thỏa đáng cho việc huấn luyện mô hình machine learning. Do đó, IterableDataset.shuffle() cũng sẽ trộn các shard của dataset nếu dataset của bạn được tạo từ nhiều tệp hoặc nguồn khác nhau. Ngoài ra, nó lấp đầy bộ đệm bằng cách sử dụng tối đa max_buffer_input_shards shard mỗi lần (mặc định là 10), điều này cải thiện đáng kể chất lượng của quá trình trộn:

# Stream from the internet
my_iterable_dataset = load_dataset("deepmind/code_contests", split="train", streaming=True)
my_iterable_dataset.num_shards  # 39

# Stream from local files
data_files = {"train": [f"path/to/data_{i}.csv" for i in range(1024)]}
my_iterable_dataset = load_dataset("csv", data_files=data_files, split="train", streaming=True)
my_iterable_dataset.num_shards  # 1024

# From a generator function
def my_generator(n, sources):
    for source in sources:
        for example_id_for_current_source in range(n):
            yield {"example_id": f"{source}_{example_id_for_current_source}"}

gen_kwargs = {"n": 10, "sources": [f"path/to/data_{i}" for i in range(1024)]}
my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs=gen_kwargs)
my_iterable_dataset.num_shards  # 1024

Nếu tập dữ liệu của bạn có ít tệp, bạn có thể cải thiện chất lượng xáo trộn bằng cách gọi reshard() trước khi xáo trộn. Ví dụ, đối với các tập dữ liệu Parquet, thao tác này sẽ chia lại shard dựa trên các nhóm dòng (row groups) thay vì mỗi shard chỉ chứa một tệp, giúp bạn có nhiều shard hơn để xáo trộn.

Chênh lệch tốc độ

Các đối tượng Dataset thông thường được xây dựng dựa trên Arrow, cung cấp khả năng truy cập ngẫu nhiên nhanh chóng vào các dòng. Nhờ vào memory mapping và việc Arrow là một định dạng trong bộ nhớ, việc đọc dữ liệu từ đĩa không thực hiện các hệ thống gọi đắt đỏ và không cần giải chuỗi hóa. Nó cung cấp tốc độ nạp dữ liệu nhanh hơn nữa khi lặp lại bằng vòng lặp for thông qua việc lặp trên các lô ghi Arrow liên tục.

Tuy nhiên, ngay khi Dataset của bạn có một ánh xạ chỉ số (ví dụ thông qua Dataset.shuffle()), tốc độ có thể chậm đi 10 lần. Điều này là do có một bước bổ sung để lấy chỉ số hàng cần đọc bằng cách sử dụng ánh xạ chỉ số, và quan trọng nhất là bạn không còn đọc các khối dữ liệu liên tục nữa. Để khôi phục tốc độ, bạn sẽ cần ghi đè lại toàn bộ tập dữ liệu trên đĩa của mình bằng cách sử dụng Dataset.flatten_indices(), hàm này sẽ loại bỏ ánh xạ chỉ số. Tuy nhiên, điều này có thể mất rất nhiều thời gian tùy thuộc vào kích thước của tập dữ liệu của bạn:

my_dataset[0]  # fast
my_dataset = my_dataset.shuffle(seed=42)
my_dataset[0]  # up to 10x slower
my_dataset = my_dataset.flatten_indices()  # rewrite the shuffled dataset on disk as contiguous chunks of data
my_dataset[0]  # fast again

Trong trường hợp này, chúng tôi khuyến nghị chuyển sang IterableDataset và tận dụng phương pháp xáo trộn gần đúng nhanh của nó IterableDataset.shuffle(). Phương pháp này chỉ xáo trộn thứ tự các shard và thêm một bộ đệm xáo trộn vào dataset của bạn, giúp duy trì tốc độ tối ưu của dataset. Bạn cũng có thể dễ dàng xáo trộn lại dataset:

for example in enumerate(my_iterable_dataset):  # fast
    pass

shuffled_iterable_dataset = my_iterable_dataset.shuffle(seed=42, buffer_size=100)

for example in enumerate(shuffled_iterable_dataset):  # as fast as before
    pass

shuffled_iterable_dataset = my_iterable_dataset.shuffle(seed=1337, buffer_size=100)  # reshuffling using another seed is instantaneous

for example in enumerate(shuffled_iterable_dataset):  # still as fast as before
    pass

Nếu bạn đang sử dụng bộ dữ liệu của mình trên nhiều epoch, hạt giống hiệu dụng để xáo trộn thứ tự các shard trong bộ đệm xáo trộn là seed + epoch. Điều này giúp việc xáo trộn lại bộ dữ liệu giữa các epoch trở nên dễ dàng:

for epoch in range(n_epochs):
    my_iterable_dataset.set_epoch(epoch)
    for example in my_iterable_dataset:  # fast + reshuffled at each epoch using `effective_seed = seed + epoch`
        pass

Để khởi động lại quá trình lặp của một tập dữ liệu dạng map, bạn chỉ cần bỏ qua các ví dụ đầu tiên:

my_dataset = my_dataset.select(range(start_index, len(dataset)))

Nhưng nếu bạn sử dụng DataLoader với một Sampler, bạn nên lưu trạng thái của sampler thay vào đó (bạn có thể đã viết một sampler tùy chỉnh cho phép tiếp tục).

Mặt khác, các tập dữ liệu dạng lặp (iterable datasets) không cung cấp khả năng truy cập ngẫu nhiên vào một chỉ mục ví dụ cụ thể để tiếp tục từ đó. Tuy nhiên, bạn có thể sử dụng IterableDataset.state_dict() và IterableDataset.load_state_dict() để tiếp tục từ một checkpoint, tương tự như cách bạn có thể làm với các mô hình và bộ tối ưu hóa:

>>> iterable_dataset = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=3)
>>> # save in the middle of training
>>> state_dict = iterable_dataset.state_dict()
>>> # and resume later
>>> iterable_dataset.load_state_dict(state_dict)

Bên trong, tập dữ liệu dạng lặp (iterable dataset) theo dõi shard hiện đang được đọc và chỉ mục của ví dụ trong shard hiện tại, sau đó lưu trữ thông tin này trong state_dict.

Để khôi phục từ một checkpoint, dataset sẽ bỏ qua tất cả các shard đã được đọc trước đó để bắt đầu lại từ shard hiện tại. Sau đó, nó đọc shard và bỏ qua các ví dụ cho đến khi đạt đến đúng ví dụ từ checkpoint.

Do đó, việc khởi động lại một dataset khá nhanh, vì nó sẽ không đọc lại các shard đã được lặp qua. Tuy nhiên, việc tiếp tục một dataset thường không diễn ra tức thì, vì nó phải bắt đầu đọc lại từ đầu shard hiện tại và bỏ qua các ví dụ cho đến khi đạt đến vị trí checkpoint.

Điều này có thể được sử dụng cùng với StatefulDataLoader từ torchdata, xem streaming với PyTorch DataLoader.

Chuyển từ kiểu map sang kiểu iterable

Nếu bạn muốn tận dụng hành vi “lazy” của IterableDataset hoặc lợi thế về tốc độ của chúng, bạn có thể chuyển Dataset theo kiểu map của mình sang IterableDataset:

my_iterable_dataset = my_dataset.to_iterable_dataset()

Nếu bạn muốn xáo trộn bộ dữ liệu hoặc dùng nó với PyTorch DataLoader, chúng tôi khuyến nghị tạo một IterableDataset dạng sharded:

my_iterable_dataset = my_dataset.to_iterable_dataset(num_shards=1024)
my_iterable_dataset.num_shards  # 1024