rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-22 03:32:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

docs(python): rework PyTorch page around StreamingDataset (#4606)

Browse Source 
## Summary
- Rework `docs/web/training/torch.mdx` for people who already know the
PyTorch `DataLoader`, built around `lancedb.streaming.StreamingDataset`.
- The quickstart trains a state-to-action policy on the real
`lance-format/lerobot-xvla-soft-fold` dataset loaded via `hf://`, using
a zero-copy Arrow-to-tensor transform.
- Add a short table mapping PyTorch idioms (`IterableDataset`,
`DistributedSampler`, `collate_fn`, `random_split`,
`StatefulDataLoader`) to their LanceDB equivalents.
- Add an Optimizations section: Transform (including an `ExtractFrames`
example using blob v2 and torchcodec), Distributed Training, and the
existing Output Formats content moved underneath.

## Notes
- The quickstart depends on #4602 (permutation building over tables with
schema metadata), which is merged.
- `ExtractFrames` takes a URI and opens the table lazily because a local
`LanceTable` cannot yet be pickled into DataLoader workers.
- The Permutation-based "Selecting columns" and "Using multiple
DataLoader workers" sections further down are unchanged from main.

## Test plan
- [x] Ran the quickstart, transform, `ExtractFrames` (main process and
forkserver workers) and distributed snippets locally.
- [x] Rendered the page with `mint dev` and checked for console errors.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Lei Xu
2026-10-20 21:47:11 -07:00
committed by GitHub
co-authored by Claude Sonnet 5.5
parent 951ae51ef6
commit bedb6ab3b1
1 changed files with 135 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/web/training/torch.mdx
+135 -28
View File
@@ -1,7 +1,7 @@
---
title: "PyTorch Integration"
sidebarTitle: "PyTorch"
description: Learn how to use LanceDB with PyTorch for training and inference.
description: Use a LanceDB table as a PyTorch IterableDataset, mapped onto the `torch.utils.data` concepts you already know.
icon: fire
availability:
oss: available
@@ -11,44 +11,151 @@ availability:
tables in the Python SDK.
---
Use LanceDB tables as input to a PyTorch `DataLoader`. Select the columns you
need, choose a batch format, and read data in parallel for training or inference.
`StreamingDataset` is a LanceDB table exposed as a PyTorch
[`IterableDataset`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset),
so it drops into a stock [`DataLoader`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader).
It handles per-epoch shuffling, sharding across ranks and workers, prefetching, and exact resume, so you
don't write those yourself. The shuffle is fixed when the dataset is constructed, so each epoch needs a new
`StreamingDataset` built with that epoch's number.
## Quickstart {#quickstart}
The `Table` class in LanceDB implements a contract for a PyTorch
[Dataset](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset).
This means you can simply use a LanceDB table in a PyTorch dataloader directly.
```py Python icon=Python
import lancedb, pyarrow as pa, torch
from torch import nn
from lancedb.streaming import StreamingDataset
```py Python icon=Python
import lancedb
import torch
import pyarrow as pa
from lancedb.util import tbl_to_tensor
def to_samples(rb: pa.RecordBatch):
f = lambda c: torch.from_dlpack(rb[c].flatten()).view(len(rb), -1)
rows = zip(f("observation_state"), f("action"))
return [{"state": s, "action": a} for s, a in rows]
mem_db = lancedb.connect("memory://")
table = mem_db.create_table("test_table", pa.table({"a": range(1000)}))
uri = "hf://datasets/lance-format/lerobot-xvla-soft-fold/data"
db = lancedb.connect(uri, storage_options={"hf_enable_resolve_cache": "true"})
table = db.open_table("frames")
train = StreamingDataset(
table, columns=["observation_state", "action"],
transform=to_samples)
# Any LanceDB table can be used as a PyTorch Dataset
dataloader = torch.utils.data.DataLoader(
table, batch_size=1024, shuffle=True, collate_fn=tbl_to_tensor
)
for batch in dataloader:
print(batch)
model = nn.Sequential(nn.Linear(96, 256), nn.ReLU(), nn.Linear(256, 14))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
for batch in torch.utils.data.DataLoader(train, batch_size=256):
loss = nn.MSELoss()(model(batch["state"]), batch["action"])
opt.zero_grad()
loss.backward()
opt.step()
```
Although the `Table` class in LanceDB implements the `torch.utils.data.Dataset` interface, you may find that using
a table [Permutation](/training/data-loading) is more flexible.
`StreamingDataset` is a standard PyTorch `IterableDataset`, so it works with everything built for one: a stock
`DataLoader`, its default `collate_fn`, and worker processes. The transform returns plain tensors, so the
model, loss and optimizer are unmodified PyTorch. Here it streams the
[`lerobot-xvla-soft-fold`](https://huggingface.co/datasets/lance-format/lerobot-xvla-soft-fold) dataset
straight from the Hugging Face Hub to train a policy that predicts each `action` from the robot's
`observation_state`.
```py Python icon=Python
from lancedb.permutation import Permutation
## PyTorch concepts in LanceDB {#concept-map}
permutation = Permutation.identity(table)
dataloader = torch.utils.data.DataLoader(permutation)
| In PyTorch you would write | In LanceDB you use |
| --- | --- |
| An `IterableDataset` that reads files in `__iter__` | `StreamingDataset(table)` |
| `DistributedSampler`, or sharding by `get_worker_info()` in `__iter__` | `rank=` and `world_size=`. Each rank and each DataLoader worker reads its own disjoint splits |
| Decoding in `__getitem__` or `collate_fn` | `transform=`, which receives a whole Arrow `RecordBatch` and runs in a thread pool |
| `random_split`, `Subset`, or reading only the fields a model needs | `filter="<SQL predicate>"` and `columns=[...]`. Unselected columns are never read from storage |
| `StatefulDataLoader` | `state_dict()` and `load_state_dict()`, plus `StreamingDataLoader` for workers |
The dataset does not copy your data. At construction it computes a shuffled assignment of row IDs to
`num_splits` fixed splits, then reads rows from the table as it iterates. It reads the table version
that exists at construction, so rows appended later do not appear mid-epoch.
## Optimizations {#optimizations}
The quickstart works as written. These sections cover what to tune as training scales up: the transform that
builds each sample, running across many GPUs, and the output format of a `Table` or `Permutation`.
### Transform {#transform}
A `transform` turns a batch of Arrow rows into the per-sample items your `DataLoader` collates. It receives a
whole `pyarrow.RecordBatch` (`read_batch_size` rows, 64 by default) rather than one sample at a time, so
vectorized code is cheap, and it must return exactly one item per row. The quickstart's transform is the
whole pattern: columns become tensors and each row becomes a dict, which the default `collate_fn` stacks into
a batch. Without a transform, rows arrive as plain Python dicts.
Heavier per-sample work belongs in the same function, such as decoding video. Store each clip in a blob v2
column (`lancedb.blob("video")`). A scan returns small descriptors instead of the video bytes, so reading
other columns never pulls video in. `fetch_blob_files` opens the clips for a batch of rows as file handles,
and torchcodec's `VideoDecoder` decodes just the frames you ask for. `ExtractFrames` samples 8 evenly spaced
frames from each clip:
```py Python icon=Python
import functools
from torchcodec.decoders import VideoDecoder
class ExtractFrames:
def __init__(self, uri: str, name: str, num_frames: int = 8) -> None:
self.uri, self.name, self.num_frames = uri, name, num_frames
@functools.cached_property
def table(self) -> lancedb.table.Table:
return lancedb.connect(self.uri).open_table(self.name)
def __call__(self, rb: pa.RecordBatch) -> list[dict[str, torch.Tensor]]:
samples = []
for blob in self.table.fetch_blob_files("video", rb["_rowid"]):
decoder = VideoDecoder(blob)
idx = torch.linspace(0, len(decoder) - 1, self.num_frames).long()
samples.append({"frames": decoder.get_frames_at(idx.tolist()).data})
return samples
# `clips` is the table "clips" in the database at `uri`, with a blob v2 column "video"
train = StreamingDataset(clips, columns=["_rowid"], transform=ExtractFrames(uri, "clips"))
```
## Output Formats {#output-formats}
Transforms run in a thread pool, one thread per CPU by default (`transform_parallelism`), overlapping with
storage reads and your training step. NumPy, PyArrow and PyTorch release the GIL, so vectorized
work runs in parallel. A pure-Python loop does not. If a profile shows the transform is the bottleneck, move
it into [DataLoader workers](#using-multiple-dataloader-workers).
Real data has bad rows, such as a truncated frame. Raise inside the transform and set
`on_transform_error="skip"` (or `"warn"` to log each failure). LanceDB re-runs the failing batch one row at a
time, drops only the rows that raise, and counts them in `dataset.rows_skipped`. When the bad rows match a SQL
predicate, prefer `filter=`: it runs before the splits are built, so every split keeps the same length.
### Distributed Training {#distributed-training}
`StreamingDataset` takes the place of `DistributedSampler`. Pass the rank and world size, and each rank reads
its own disjoint part of the table, together covering it once per epoch:
```py Python icon=Python
import torch.distributed as dist
dist.init_process_group("nccl")
train = StreamingDataset(
table, columns=["observation_state", "action"], transform=to_samples,
num_splits=48, shuffle_seed=42, epoch=epoch,
rank=dist.get_rank(), world_size=dist.get_world_size())
```
Everything after the dataset is standard distributed PyTorch. Wrap the model in `DistributedDataParallel`
and pass `batch_size` to the `DataLoader` per GPU, so the global batch is `batch_size * world_size`. Nothing
in the data path changes.
Three arguments matter in distributed runs:
- **`num_splits`** is how many fixed streams the table is divided into. It must be divisible by `world_size`,
and by `world_size * num_workers` when you use DataLoader workers. Choose a number with many divisors, such
as 48, and runs on 8, 4 or 2 GPUs see the same data at every step, provided the global batch size is a multiple of
`num_splits`.
- **`shuffle_seed` and `epoch`** must be identical on every rank. Pass a fixed `shuffle_seed` and increase
`epoch` each epoch. `shuffle_seed=None` draws a different random seed in each process, so ranks would read
overlapping rows and miss others.
- **`rank` and `world_size`** come from your launcher, for example `torchrun`.
To resume after a failure, save `train.state_dict()` next to the model. It records progress per split, so the
resumed run continues from the same row even on a different number of GPUs. See
[Parallelism & Elasticity](/training/data-loading#parallelism-&-elasticity) and
[Checkpointing](/training/data-loading#checkpointing) for the batch-size rules and the multi-worker case.
### Output Formats {#output-formats}
By default, a `Table` data loader will emit Arrow data. `collate_fn` is PyTorch's batching hook: PyTorch calls it to
turn the fetched items into one batch. PyTorch's default collate function only knows how to combine tensors, NumPy
@@ -68,7 +175,7 @@ the Arrow data in different ways. The `arrow` and `polars` formats will always
`pandas`, and `torch_col` formats will also avoid data copies in most cases. The `python`, `python_col`, and
`torch` formats will all require at least one full copy of the data and are the slowest options.
### Using the torch column format with a torch data loader {#using-the-torch-column-format-with-a-torch-data-loader}
#### Using the torch column format with a torch data loader {#using-the-torch-column-format-with-a-torch-data-loader}
The `torch_col` format is the most efficient way to convert from Arrow to a `torch.Tensor`. It will convert the
entire Arrow batch to a _column-major_ `torch.Tensor`. In other words, given C columns and R rows, the resulting