Lightning-AI/litData

LitData – PyTorch用の高速でクラウドネイティブなデータパイプライン

何であるか – LitData（Lightning AI製）は、巨大なトレーニングデータセットを高速かつ信頼性高く読み込むためのPythonライブラリです。これは、クラウドストレージ（S3、GCS、Azure、Hugging Face Hubなど）から直接生のファイルをストリーミングするか、一度だけデータをコンパクトなチャンク化バイナリ形式に変換することで、単純な torch.utils.data.DataLoader よりも最大20倍の速度でストリーミングできるようにします。

なぜ重要か – 現代のディープラーニングプロジェクトでは、I/Oがしばしばボトルネックになります。データセットはテラバイト規模であり、数百万の小さなファイルで構成され、またはリモートのオブジェクトストアにのみ存在する場合があります。LitDataは以下の3つの課題を解決します：

ゼロ準備ストリーミング – StreamingRawDataset は、完全に非同期でバッチダウンロードを行い、組み込みの再試行ロジックを備えています。ローカルにデータをすべてダウンロードする必要はありません。
一度だけの最適化 – ld.optimize はデータセットをLitData独自のチャンク化バイナリ形式（またはParquet、MDSなども対応）に再書き込みします。生成されたチャンクはローカルおよびバケットにキャッシュされ、再開可能でシャッフル対応のストリーミングが可能になり、最大20倍の高速化が実現します。
分散対応API – ライブラリには StreamingDataset / StreamingDataLoader（PyTorch Lightning、Fabric、Hugging Faceと互換）が同梱されており、map プリミティブにより、画像のリサイズ、埋め込みの生成、ウェブスクレイピングなど、任意の前処理を多数のマシンで並列実行できます。

コアコンセプトとAPI

コンセプト	代表的なクラス/関数	機能
生ストリーミング	`StreamingRawDataset`	ローカルまたはクラウドパスから任意のファイル（画像、音声、テキストなど）を生の `bytes` として読み込みます。`transform` コールバックを指定することで、実行時にデコードできます。
最適化ストリーミング	`StreamingDataset` + `StreamingDataLoader`	`ld.optimize` で事前に変換されたデータを読み込みます。シャッフル、drop-last、マルチGPU再開、キーによる検索を処理します。
一度だけの変換	`ld.optimize(...)`	サンプルを生成するPython関数を受け取り、LitDataのチャンク化バイナリ形式（チャンクサイズは設定可能）に書き込み、オプションでランダムアクセス用のキーインデックスを構築します。
並列前処理	`ld.map(...)`	ユーザー提供の関数を入力リストに対して並列実行し、結果を出力先（ローカルまたはクラウド）に書き込みます。画像のリサイズ、埋め込み生成、ウェブスクレイピングなどに最適です。
Hugging Face統合	`ld.optimize_hf(...)` / `StreamingDataset("hf://…")`	HF Hubから直接データセットをストリーミングするか、一度だけLitDataチャンクに変換して高速トレーニングを可能にします。

クイックスタート（インストールと基本使用）

pip install litdata                # コアパッケージ
pip install 'litdata[extras]'      # より高速な asyncio 用に uvloop を追加

1️⃣ 生ファイルのストリーミング（前処理ステップ不要）

from litdata import StreamingRawDataset
from torch.utils.data import DataLoader
from PIL import Image
import io

ds = StreamingRawDataset(
    "s3://my-bucket/raw-images/",
    transform=lambda b: Image.open(io.BytesIO(b)).convert("RGB"),
)
loader = DataLoader(ds, batch_size=32, num_workers=8)
for batch in loader:
    train_step(batch)

特徴：非同期バッチダウンロード、自動再試行、ローカル index.json.zstd キャッシュ、あらゆるクラウドプロバイダに対応。

2️⃣ 一度最適化、その後最大速度でストリーミング

import litdata as ld, numpy as np

def make_sample(i):
    img = np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8)
    return {"index": i, "image": ld.Image(array=img, quality=95, format="jpeg"), "label": np.random.randint(10)}

ld.optimize(fn=make_sample, inputs=list(range(1000)), output_dir="fast_data", chunk_bytes="64MB")
# フォルダをクラウドにアップロード、例：
# aws s3 cp --recursive fast_data s3://my-bucket/fast_data

ds = ld.StreamingDataset('s3://my-bucket/fast_data', shuffle=True, seed=42)
loader = ld.StreamingDataLoader(ds, batch_size=64)
for batch in loader:
    # batch["image"] は Image オブジェクトのリスト、batch["label"] は整数のリスト
    train_step(batch)

結果：同じ生ファイルに対して、従来の torch.utils.data.DataLoader と比較して最大20倍の高速なエポック処理時間。

LitData vs. 代替手段の使い分け

LitData – クラウド非依存ストリーミング、再開可能なエポック、大規模シャッフル、一度の前処理で複数実験に再利用したい場合に最適。
torchdata – 低レベルのイテレータプリミティブを提供しますが、組み込みのストレージ形式やチャンクレベルのキャッシュはなし。シンプルなファイル一覧が必要で、他のすべての部分を自分で構築したい場合に使用。
Hugging Face datasets ストリーミング – 早期プロトタイピングに便利。LitDataの optimize_hf は、同じデータセットを長期間のトレーニングに向け、さらに高速化できます。

対応するデータモダリティ

LitDataは、書き込み/読み込み時に型情報を保持する軽量ラッパーを提供しています：

テキスト / トークン – Text, Tokens
画像 / JPEG / PIL – Image, Jpeg, Pil
音声 / 動画 – Audio, Video
3Dメッシュ、Niftiボリューム – Mesh, Nifti
汎用ファイル – File, Pdf
配列とテンソル – Numpy, Tensor
グラフ（PyG） – Graph
Parquetテーブル – Parquet これらのラッパーにより、生のバイト、NumPy配列、PyTorchテンソルを保存でき、トレーニング時に同じPythonオブジェクトとして取得できます。

エコシステムとコミュニティ

Lightning Cloud – シームレスな統合。LightningのマネージドGPUクラスタやオンプレミスマシンでLitDataパイプラインを実行可能。
Discord – 活発なサポートチャンネル（https://discord.gg/VptPCZkGNa）。
AIエージェントスキル – Vercelスタイルの「スキル」が提供され、コード生成エージェント（Claude、Cursorなど）がLitData APIの補完を可能にします。
使用実績 – Lightning Cloud上で34万以上の開発者、内部研究チーム、複数の公開MLプロジェクトで利用（リポジトリの「Used by」セクションにリンクあり）。

TL;DR

LitDataは、PyTorch向けのプロダクションレベルのデータ読み込みライブラリで、以下の機能を提供します：

任意のクラウドストアから生ファイルを非同期バッチI/Oでストリーミング。
一度だけデータセットを高速なチャンク化バイナリ形式に変換し、シャッフル、再開、キーによるアクセスが可能。
ld.map を使ってマシン間で任意の前処理を並列化。
Drop-inクラスである StreamingDataset / StreamingDataLoader で、Lightning、Hugging Face、PyTorch-Lightningワークフローに統合。

トレーニングジョブがデータ待ちに多くの時間を費やしている場合、LitDataはその待機時間を大幅に短縮し、大規模なクラウドベースデータセットのエンジニアリングを簡素化します。