Crate burn_store 

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
• bridge: Conversions between burn-core tensors and the burn_pack::Tensor entries the stores move around

Feature Flags

• std: Enables file I/O and other std-only features (default)
• safetensors: Enables SafeTensors format support (default)
• pytorch: Enables loading PyTorch .pt/.pth files (default)
• memmap: Implied by std and adds nothing on top of it. Loading a SafeTensors file memory-maps it whenever std is on, regardless of this flag.

Re-exports

pub use pytorch::PytorchStore;

pub use pytorch::PytorchStoreError;

pub use pytorch_reader;

pub use burn_pack;

Modules

bridge

Bridge between burn-core tensors and the tensor-agnostic burnpack format entries.

pytorch

PyTorch format support for burn-store.

Structs

Applier

Applier that applies loaded tensors to module parameters with proper adapter support using container type information

ApplyResult

Result of applying tensors 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

ChainAdapter

Adapter that applies two adapters in sequence.

Collector

Collects a module’s tensors without copying data.

FloatCastAdapter

Adapter that casts every float tensor to one target dtype.

HalfPrecisionAdapter

Adapter for mixed-precision (F32/F16) model storage.

IdentityAdapter

Identity adapter that passes tensors through unchanged

KeyRemapper

Key remapper for transforming tensor names.

ModuleContext

Where in a module hierarchy a tensor was found.

PathFilter

A sophisticated path filter that supports multiple matching strategies.

PyTorchToBurnAdapter

Adapter for converting from PyTorch format to Burn format

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.

Traits

ModuleAdapter

Trait for adapting tensors 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.

map_indices_contiguous_except

Like map_indices_contiguous, but leaves the indices under any prefix for which keep returns true untouched.