Lightning-AI/litData

LitData – 面向 PyTorch 的快速、雲原生資料管道

是什麼 – LitData（由 Lightning AI 開發）是一個 Python 庫，可讓載入大型訓練資料集變得快速且可靠。它透過兩種方式實現：一是直接從雲端儲存（S3、GCS、Azure、Hugging Face Hub 等）串流傳輸原始檔案；二是將資料一次性轉換為緊湊的分塊二進位格式，可實現比原始 torch.utils.data.DataLoader 快達 20 倍的串流傳輸、洗牌與可恢復的 epoch。

為何重要 – 在現代深度學習專案中，I/O 通常是瓶頸：資料集可能達到 TB 級別，包含數百萬個小型檔案，或僅存在於遠端物件儲存中。LitData 解決了三個痛點：

零準備串流傳輸 – 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]'      # 加入可選的 uvloop 以提升 asyncio 速度

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 相比，epoch 時間最多可提升 20 倍。

何時使用 LitData 與替代方案

LitData – 當你需要雲端無關串流傳輸、可恢復 epoch、大規模洗牌，或希望一次性預處理並跨多個實驗複用優化格式時最佳。
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。
使用情況 – 超過 34 萬名開發者在 Lightning Cloud 上使用，內部研究團隊及多個公開 ML 專案（詳見倉儲「Used by」部分連結）。

TL;DR

LitData 是一個生產級 PyTorch 資料載入庫，可實現：

直接從任意雲端儲存串流傳輸原始檔案，支援非同步批次 I/O。
一次性將資料集轉換為快速、分塊的二進位格式，支援洗牌、恢復與基於鍵的存取。
透過 ld.map 在多台機器上並行化任意預處理。
透過 StreamingDataset/StreamingDataLoader 類無縫整合 Lightning、Hugging Face 和 PyTorch-Lightning 工作流程。

如果你的訓練任務花費大量時間等待資料，LitData 可大幅縮短等待時間，並簡化大規模雲端資料集的工程實作。