Crate burn_store

Crate burn_store 

Source
Expand description

§Burn Store

Advanced model storage and serialization infrastructure for the Burn deep learning framework.

This crate provides comprehensive functionality for storing and loading Burn modules and their tensor data, with support for cross-framework interoperability, flexible filtering, and efficient memory management through lazy materialization.

§Key Features

  • Burnpack Format: Native Burn format with CBOR metadata, ParamId persistence for stateful training, and no-std support
  • SafeTensors Format: Industry-standard format for secure and efficient tensor serialization
  • PyTorch Compatibility: Load PyTorch models directly into Burn with automatic weight transformation
  • Zero-Copy Loading: Memory-mapped files and lazy tensor materialization for optimal performance
  • Flexible Filtering: Load/save specific model subsets using regex, exact paths, or custom predicates
  • Tensor Remapping: Rename tensors during load/save operations for framework compatibility
  • No-std Support: Core functionality available in embedded and WASM environments

§Quick Start

§Basic Save and Load

ⓘ
use burn_store::{ModuleSnapshot, SafetensorsStore};

// Save a model
let mut store = SafetensorsStore::from_file("model.safetensors");
model.save_into(&mut store)?;

// Load a model
let mut store = SafetensorsStore::from_file("model.safetensors");
model.load_from(&mut store)?;

§Loading PyTorch Models

ⓘ
use burn_store::PytorchStore;

// Load PyTorch model (automatic weight transformation via PyTorchToBurnAdapter)
let mut store = PytorchStore::from_file("pytorch_model.pth")
    .with_top_level_key("state_dict")  // Access nested state dict if needed
    .allow_partial(true);               // Skip unknown tensors

model.load_from(&mut store)?;

§Filtering and Remapping

// Save only specific layers with renaming
let mut store = SafetensorsStore::from_file("encoder.safetensors")
    .with_regex(r"^encoder\..*")                         // Filter: only encoder layers
    .with_key_remapping(r"^encoder\.", "transformer.")   // Rename: encoder.X -> transformer.X
    .metadata("subset", "encoder_only");

// Use store with model.save_into(&mut store)?;

§Core Components

  • ModuleSnapshot: Extension trait for Burn modules providing collect() and apply() methods
  • BurnpackStore: Native Burn format with ParamId persistence for stateful training workflows
  • SafetensorsStore: Primary storage implementation supporting the SafeTensors format
  • PytorchStore: PyTorch model loader supporting .pth and .pt files
  • PathFilter: Flexible filtering system for selective tensor loading/saving
  • KeyRemapper: Advanced tensor name remapping with regex patterns
  • ModuleAdapter: Framework adapters for cross-framework compatibility

§Feature Flags

  • std: Enables file I/O and other std-only features (default)
  • safetensors: Enables SafeTensors format support (default)

Re-exports§

pub use pytorch::PytorchStore;
pub use pytorch::PytorchStoreError;

Modules§

pytorch
PyTorch format support for burn-store.

Structs§

Applier
Applier that applies tensor snapshots to module parameters with proper adapter support using container type information
ApplyResult
Result of applying tensor snapshots to a module
BurnToPyTorchAdapter
Adapter for converting from Burn format to PyTorch format
BurnpackStore
BurnpackStore - A Burn-specific file format store using CBOR for metadata
BurnpackWriter
Writer for creating Burnpack files
Collector
Collects tensor views from modules without copying data.
IdentityAdapter
Identity adapter that passes tensors through unchanged
KeyRemapper
Key remapper for transforming tensor names.
PathFilter
A sophisticated path filter that supports multiple matching strategies.
PyTorchToBurnAdapter
Adapter for converting from PyTorch format to Burn format
TensorSnapshot
A lightweight snapshot of a tensor that can lazily produce TensorData.

Enums§

ApplyError
Error types that can occur during tensor application
SafetensorsStore
SafeTensors store supporting both file and memory storage.
SafetensorsStoreError
Errors that can occur during SafeTensors operations.
TensorSnapshotError
Error type for TensorSnapshot operations

Traits§

ModuleAdapter
Trait for adapting tensor snapshots between different module formats
ModuleSnapshot
Extension trait for modules that provides tensor storage functionality.
ModuleStore
A trait for handling module storage operations.

Functions§

map_indices_contiguous
Map tensor paths to have contiguous numeric indices.