Perbedaan antara Dataset dan IterableDataset

Terdapat dua jenis objek dataset, yaitu Dataset dan IterableDataset. Jenis dataset yang Anda pilih untuk digunakan atau dibuat bergantung pada ukuran dataset tersebut. Secara umum, IterableDataset sangat ideal untuk dataset besar (bayangkan ratusan GB!) berkat perilaku lazy-nya dan keunggulan kecepatannya, sedangkan Dataset sangat cocok untuk semua kasus lainnya. Halaman ini akan membandingkan perbedaan antara Dataset dan IterableDataset untuk membantu Anda memilih objek dataset yang tepat.

Mengunduh dan streaming

Ketika Anda memiliki Dataset reguler, Anda dapat mengaksesnya menggunakan my_dataset[0]. Ini menyediakan akses acak ke baris-barisnya. Dataset seperti ini juga disebut dataset “map-style”. Sebagai contoh, Anda dapat mengunduh ImageNet-1k seperti ini dan mengakses baris mana pun:

from datasets import load_dataset

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

Namun, satu hal yang perlu diperhatikan adalah Anda harus menyimpan seluruh dataset di disk atau di memori, yang menghalangi Anda dari mengakses dataset yang lebih besar dari disk. Karena hal ini bisa menjadi tidak praktis untuk dataset besar, ada jenis dataset lain, yaitu IterableDataset. Ketika Anda memiliki IterableDataset, Anda dapat mengaksesnya menggunakan loop for untuk memuat data secara bertahap saat Anda mengiterasi dataset. Dengan cara ini, hanya sebagian kecil contoh yang dimuat ke dalam memori, dan Anda tidak menulis apa pun ke disk.

Sebagai contoh, Anda dapat melakukan streaming dataset ImageNet-1k tanpa mengunduhnya ke disk:

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 dapat membaca data online tanpa menulis file apa pun ke disk. Sebagai contoh, Anda dapat melakukan streaming pada dataset yang terdiri dari beberapa shard, masing-masing berukuran ratusan gigabyte seperti C4 atau LAION-2B. Pelajari lebih lanjut tentang cara melakukan streaming dataset di Panduan Streaming Dataset.

Namun, ini bukan satu-satunya perbedaan, karena perilaku “lazy” dari IterableDataset juga hadir saat pembuatan dan pemrosesan dataset.

Membuat dataset bergaya peta dan dataset yang dapat diiterasi

Anda dapat membuat Dataset menggunakan daftar atau kamus, dan data sepenuhnya dikonversi ke Arrow sehingga Anda dapat dengan mudah mengakses baris mana pun:

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

Di sisi lain, untuk membuat IterableDataset, Anda harus menyediakan cara “lazy” untuk memuat data. Di Python, kita umumnya menggunakan fungsi generator. Fungsi-fungsi ini yield satu contoh pada satu waktu, yang berarti Anda tidak dapat mengakses baris dengan memotongnya seperti Dataset biasa:

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

Memuat file lokal secara keseluruhan dan bertahap

Anda dapat mengonversi file data lokal atau jarak jauh menjadi Dataset Arrow menggunakan load_dataset():

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

Namun, hal ini memerlukan langkah konversi dari format CSV ke Arrow, yang memakan waktu dan ruang disk jika dataset Anda besar.

Untuk menghemat ruang disk dan melewati langkah konversi, Anda dapat mendefinisikan IterableDataset dengan melakukan streaming langsung dari file lokal. Dengan cara ini, data dibaca secara progresif dari file lokal saat Anda mengiterasi dataset:

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

Banyak format file yang didukung, seperti CSV, JSONL, dan Parquet, serta file gambar dan audio. Anda dapat menemukan informasi lebih lanjut di panduan yang sesuai untuk memuat dataset tabular, teks, visi, dan audio.

Pemrosesan data eager dan pemrosesan data lazy

Saat Anda memproses objek Dataset menggunakan Dataset.map(), seluruh dataset diproses secara langsung dan dikembalikan. Ini mirip dengan cara kerja pandas, misalnya.

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

Di sisi lain, karena sifat “lazy” dari IterableDataset, pemanggilan IterableDataset.map() tidak menerapkan fungsi map Anda pada seluruh dataset. Sebagai gantinya, fungsi map Anda diterapkan secara on-the-fly.

Karena hal itu, Anda dapat merangkai beberapa langkah pemrosesan dan semuanya akan berjalan sekaligus saat Anda mulai mengiterasi dataset:

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

Pengacakan eksak dan aproksimasi cepat

Saat Anda mengacak Dataset menggunakan Dataset.shuffle(), Anda menerapkan pengacakan yang tepat pada dataset. Cara kerjanya adalah dengan mengambil daftar indeks [0, 1, 2, ... len(my_dataset) - 1] dan mengacak daftar tersebut. Kemudian, mengakses my_dataset[0] akan mengembalikan baris dan indeks yang didefinisikan oleh elemen pertama dari pemetaan indeks yang telah diacak:

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

Karena kita tidak memiliki akses acak ke baris dalam kasus IterableDataset, kita tidak dapat menggunakan daftar indeks yang diacak dan mengakses baris pada posisi arbitrer. Hal ini mencegah penggunaan pengacakan yang tepat. Sebagai gantinya, pengacakan aproksimasi yang cepat digunakan dalam IterableDataset.shuffle(). Metode ini menggunakan buffer pengacakan untuk mengambil contoh acak secara iteratif dari dataset. Karena dataset masih dibaca secara iteratif, metode ini memberikan kinerja kecepatan yang sangat baik:

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

Namun, penggunaan shuffle buffer saja tidak cukup untuk menyediakan pengacakan yang memuaskan untuk pelatihan model machine learning. Oleh karena itu, IterableDataset.shuffle() juga mengacak shard dataset jika dataset Anda terdiri dari beberapa file atau sumber. Selain itu, buffer diisi menggunakan hingga max_buffer_input_shards shard pada satu waktu (default adalah 10), yang sangat meningkatkan kualitas pengacakan:

# 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

Jika dataset Anda memiliki sedikit file, Anda dapat meningkatkan kualitas pengacakan dengan memanggil reshard() sebelum pengacakan. Misalnya, untuk dataset Parquet, ini melakukan resharing menggunakan grup baris alih-alih memiliki satu file per shard, sehingga memberikan Anda banyak shard untuk diacak.

Perbedaan kecepatan

Objek Dataset reguler didasarkan pada Arrow yang menyediakan akses acak cepat ke baris. Berkat memory mapping dan fakta bahwa Arrow adalah format in-memory, membaca data dari disk tidak melakukan panggilan sistem dan deserialisasi yang mahal. Ini menyediakan pemuatan data yang bahkan lebih cepat saat iterasi menggunakan loop for dengan mengiterasi batch record Arrow yang berurutan.

Namun begitu Dataset Anda memiliki pemetaan indeks (misalnya melalui Dataset.shuffle()), kecepatannya bisa menjadi 10 kali lebih lambat. Hal ini terjadi karena ada langkah tambahan untuk mendapatkan indeks baris yang akan dibaca menggunakan pemetaan indeks, dan yang paling penting, Anda tidak lagi membaca potongan data yang berurutan. Untuk memulihkan kecepatan, Anda perlu menulis ulang seluruh dataset ke disk Anda menggunakan Dataset.flatten_indices(), yang menghapus pemetaan indeks. Namun, hal ini mungkin memakan waktu lama tergantung pada ukuran dataset Anda:

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

Dalam kasus ini, kami merekomendasikan untuk beralih ke IterableDataset dan memanfaatkan metode pengacakan perkiraan yang cepat IterableDataset.shuffle(). Metode ini hanya mengacak urutan shard dan menambahkan buffer pengacakan ke dataset Anda, yang menjaga kecepatan dataset Anda tetap optimal. Anda juga dapat mengacak ulang dataset dengan mudah:

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

Jika Anda menggunakan dataset Anda pada beberapa epoch, seed efektif untuk mengacak urutan shard dalam buffer shuffle adalah seed + epoch. Hal ini memudahkan pengacakan ulang dataset antar epoch:

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

Untuk memulai ulang iterasi dataset bergaya peta, Anda cukup melewati contoh-contoh pertama:

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

Namun jika Anda menggunakan DataLoader dengan Sampler, Anda sebaiknya menyimpan state sampler Anda (Anda mungkin telah menulis sampler kustom yang memungkinkan resume).

Di sisi lain, dataset iterable tidak menyediakan akses acak ke indeks contoh tertentu untuk melanjutkan dari sana. Namun, Anda dapat menggunakan IterableDataset.state_dict() dan IterableDataset.load_state_dict() untuk melanjutkan dari checkpoint, serupa dengan yang dapat Anda lakukan untuk model dan optimizer:

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

Di balik layar, dataset iterable melacak shard saat ini yang sedang dibaca dan indeks contoh di shard saat ini, lalu menyimpan informasi ini di state_dict.

Untuk melanjutkan dari checkpoint, dataset melewati semua shard yang sebelumnya dibaca untuk memulai ulang dari shard saat ini. Kemudian, dataset membaca shard tersebut dan melewati contoh-contoh hingga mencapai contoh yang tepat dari checkpoint.

Oleh karena itu, memulai ulang dataset cukup cepat karena tidak akan membaca ulang shard yang sudah diiterasi. Namun, melanjutkan dataset umumnya tidak instan karena harus memulai ulang pembacaan dari awal shard saat ini dan melewati contoh-contoh hingga mencapai lokasi checkpoint.

Hal ini dapat digunakan dengan StatefulDataLoader dari torchdata, lihat streaming dengan PyTorch DataLoader.

Beralih dari gaya peta ke yang dapat diiterasi

Jika Anda ingin memanfaatkan perilaku “lazy” dari IterableDataset atau keunggulan kecepatannya, Anda dapat mengubah Dataset gaya map Anda menjadi IterableDataset:

my_iterable_dataset = my_dataset.to_iterable_dataset()

Jika Anda ingin mengacak dataset Anda atau memakainya dengan PyTorch DataLoader, kami merekomendasikan untuk menghasilkan IterableDataset yang terbagi (sharded):

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