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 可显著缩短等待时间，并简化大规模云数据集的工程实现。