WebDataset

WebDataset é uma biblioteca para escrever pipelines de E/S para grandes conjuntos de dados. Suas funcionalidades de E/S sequencial e sharding a tornam especialmente útil para transmitir grandes conjuntos de dados em streaming para um DataLoader.

O formato WebDataset

Um arquivo WebDataset é um arquivo TAR que contém uma série de arquivos de dados. Todos os arquivos de dados sucessivos com o mesmo prefixo são considerados parte do mesmo exemplo (por exemplo, um arquivo de imagem/áudio e seu rótulo ou metadados):

Os rótulos e metadados podem estar em um arquivo .json, em um .txt (para uma legenda, uma descrição) ou em um .cls (para um índice de classe).

Um WebDataset em larga escala é composto por muitos arquivos chamados shards, em que cada shard é um arquivo TAR. Cada shard tem frequentemente ~1GB, mas o conjunto de dados completo pode ter vários terabytes!

Suporte multimodal

WebDataset é projetado para conjuntos de dados multimodais, ou seja, para conjuntos de dados de imagem, áudio e/ou vídeo.

De fato, como os arquivos de mídia tendem a ser bastante grandes, a I/O sequencial do WebDataset permite leituras grandes e bufferização, resultando na melhor velocidade de carregamento de dados.

Aqui está uma lista não exaustiva dos formatos de dados suportados:

• imagem: jpeg, png, tiff
• áudio: mp3, m4a, wav, flac
• vídeo: mp4, mov, avi
• outros: npy, npz

A lista completa evolui ao longo do tempo e depende da implementação. Por exemplo, você pode encontrar quais formatos o pacote webdataset suporta no código-fonte aqui.

Streaming

O streaming de arquivos TAR é rápido porque lê blocos contíguos de dados. Ele pode ser ordens de magnitude mais rápido do que ler arquivos de dados separados um por um.

O streaming do WebDataset oferece alto desempenho tanto ao ler do disco quanto do armazenamento em nuvem, o que o torna um formato ideal para alimentar um DataLoader:

Por exemplo, aqui está como transmitir o conjunto de dados timm/imagenet-12k-wds diretamente do Hugging Face:

Primeiro, você precisa fazer login com sua conta do Hugging Face, por exemplo, usando:

hf auth login

E então você pode transmitir o conjunto de dados com WebDataset:

>>> import webdataset as wds
>>> from huggingface_hub import get_token
>>> from torch.utils.data import DataLoader

>>> hf_token = get_token()
>>> url = "https://huggingface.co/datasets/timm/imagenet-12k-wds/resolve/main/imagenet12k-train-{{0000..1023}}.tar"
>>> url = f"pipe:curl -s -L {url} -H 'Authorization:Bearer {hf_token}'"
>>> dataset = wds.WebDataset(url).decode()
>>> dataloader = DataLoader(dataset, batch_size=64, num_workers=4)

Embaralhar

Em geral, conjuntos de dados em formatos WebDataset já estão embaralhados e prontos para alimentar um DataLoader. Mas você ainda pode reembaralhar os dados com o embaralhamento aproximado do WebDataset.

Além de embaralhar a lista de shards, o WebDataset usa um buffer para embaralhar um dataset sem qualquer custo em velocidade:

Para embaralhar uma lista de arquivos fragmentados e amostrar aleatoriamente do buffer de embaralhamento:

>>> buffer_size = 1000
>>> dataset = (
...     wds.WebDataset(url, shardshuffle=True)
...     .shuffle(buffer_size)
...     .decode()
... )