torch-vault

Zero-copy PyTorch checkpoint management and SQL-driven dataset streaming built with Rust, PyO3, SQLite, and DuckDB.

Features

• Safe Checkpointing: Replaces pickle-based torch.save with safetensors memory-mapping backed by an atomic SQLite WAL index.
• Selective Weight Saving: vault.save(model, metrics=..., save_if_better="val_loss") continuously records metric history into SQLite while writing .safetensors weight files to disk only when metrics improve.
• Multi-Experiment Tracking: Tag runs with run_id (e.g. Vault("./checkpoints", run_id="bilstm_v1")) to compare multiple architectures within the same vault directory.
• SQL Dataset Streaming: Segment("data.parquet").where("quality > 0.8") or Segment("data.csv") runs DuckDB predicate pushdown on disk before streaming records into PyTorch DataLoaders.
• Hugging Face Integration: Stream Hugging Face datasets directly with Segment.from_hf("name", split=["train", "test"]) supporting arbitrary split combinations, slicing (train[:80%]), and split dictionaries.
• Zero-Copy Memory Interchange: Uses the Arrow C Stream Interface (__arrow_c_stream__ PyCapsule) to pass record batches directly into memory.
• CLI & WebApp Inspector: Includes torch-vault list and torch-vault ui for inspecting checkpoints and charting metrics in a local browser interface.

Installation

pip install torch-vault

For development build:

git clone https://github.com/ethqnol/torch-vault.git
cd torch-vault
pip install maturin
maturin develop --release

Quickstart

Checkpoint Management (Vault)

import torch
import torch.nn as nn
from torch_vault import Vault

model = nn.Linear(10, 2)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Track experiment run with overwrite protection
vault = Vault("./checkpoints", run_id="bilstm_v1", overwrite=True)

for epoch in range(1, 11):
    train_loss = 0.5 # compute training loss
    val_loss = 0.4   # compute validation loss

    # Log metrics to SQLite; save weight file on disk only if val_loss improves
    vault.save(
        model=model,
        optimizer=optimizer,
        metrics={"train_loss": train_loss, "val_loss": val_loss},
        epoch=epoch,
        save_if_better="val_loss",
        mode="min",
    )

# Load best model weights automatically by metric
best_weights = vault.load(metric="val_loss", mode="min")
model.load_state_dict(best_weights["model"])

# Export SQLite metrics history to a pandas DataFrame
df = vault.to_dataframe()

SQL Dataset Streaming (Segment)

from torch_vault import Segment
from torch_vault.transforms import Compose, SelectKeys, RenameKey, StandardScaler

pipeline = Compose([
    RenameKey(src="score", dst="quality"),
    StandardScaler(keys=["quality"]),
    SelectKeys("quality", "label"),
])

# Stream Parquet or CSV data with SQL pushdown and a 80%/10%/10% split
train_loader, val_loader, test_loader = (
    Segment("data/dataset.parquet")
    .where("quality > 0.7")
    .transform(pipeline)
    .to_dataloader(batch_size=32, split=(0.8, 0.1, 0.1), seed=42)
)

# Stream Hugging Face datasets directly with split combination:
hf_loader = (
    Segment.from_hf("stanfordnlp/imdb", split=["train", "test"])
    .where("length(text) > 50")
    .to_dataloader(batch_size=32)
)

CLI & WebApp Inspector

# Print terminal ASCII table of logged checkpoints
torch-vault list ./checkpoints

# Launch browser WebApp inspector at http://127.0.0.1:8000/
torch-vault ui ./checkpoints --port 8000

Documentation

Full interactive API reference and architecture specifications are available at ethqnol.github.io/torch-vault.

License

MIT License.