# Burn Store
> Advanced model storage and serialization for the Burn deep learning framework
[](https://crates.io/crates/burn-store)
A comprehensive storage library for Burn that enables efficient model serialization, cross-framework
interoperability, and advanced tensor management.
## Features
### Core Capabilities
- **Burnpack Format** - Native Burn format with CBOR metadata, memory-mapped loading, ParamId
persistence for stateful training, and no-std support
- **SafeTensors Format** - Industry-standard format for secure and efficient tensor serialization
- **PyTorch Support** - Direct loading of PyTorch .pth/.pt files with automatic weight
transformation
- **Zero-Copy Loading** - Memory-mapped files and lazy tensor materialization for optimal
performance
- **Cross-Framework Support** - Seamless PyTorch ↔ Burn model conversion with automatic adaptations
- **Flexible Filtering** - Load/save specific model subsets with regex, exact paths, or custom
predicates
- **Tensor Remapping** - Rename tensors during load/save for framework compatibility
- **No-std Support** - Burnpack and SafeTensors formats available in embedded and WASM environments
### Advanced Features
- **Framework Adapters** - Automatic weight transposition and parameter renaming for PyTorch
compatibility
- **Lazy Transformations** - Chain tensor transformations without materializing intermediate data
- **Partial Loading** - Continue loading even when some tensors are missing
- **Custom Metadata** - Attach version info, training details, or other metadata to saved models
## Quick Start
### Basic Save and Load
#### Burnpack (Native Format)
```rust
use burn_store::{ModuleSnapshot, BurnpackStore};
// Save a model with metadata
let mut store = BurnpackStore::from_file("model.bpk")
.metadata("version", "1.0")
.metadata("description", "My trained model");
model.save_into(&mut store)?;
// Load a model (automatically memory-mapped when available)
let mut store = BurnpackStore::from_file("model.bpk");
model.load_from(&mut store)?;
```
**Performance**: Burnpack provides faster loading times and reduced memory overhead compared to
other formats.
**Training State Persistence**: Burnpack automatically preserves parameter identifiers (ParamId) for
stateful training continuation.
#### SafeTensors
```rust
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)?;
```
### Filtering Tensors
```rust
// Save only encoder layers
let mut store = SafetensorsStore::from_file("encoder.safetensors")
.with_regex(r"^encoder\..*")
.metadata("subset", "encoder_only");
model.save_into(&mut store)?;
// Load with multiple filter patterns (OR logic)
let mut store = SafetensorsStore::from_file("model.safetensors")
.with_regex(r"^encoder\..*") // Include encoder tensors
.with_regex(r".*\.bias$") // OR include any bias tensors
.with_full_path("decoder.scale"); // OR include specific tensor
model.load_from(&mut store)?;
```
### PyTorch Interoperability
```rust
use burn_store::{PyTorchToBurnAdapter, BurnToPyTorchAdapter, PytorchStore};
// Load PyTorch .pth file directly (PyTorchToBurnAdapter is applied automatically)
// Note: skip_enum_variants defaults to true for PytorchStore
let mut store = PytorchStore::from_file("pytorch_model.pth")
.with_top_level_key("state_dict") // Access nested state dict
.allow_partial(true); // Skip unknown tensors
burn_model.load_from(&mut store)?;
// Load PyTorch model from SafeTensors
let mut store = SafetensorsStore::from_file("pytorch_model.safetensors")
.with_from_adapter(PyTorchToBurnAdapter) // Auto-transpose linear weights
.skip_enum_variants(true) // Handle enum variant name differences
.allow_partial(true); // Skip unknown PyTorch tensors
burn_model.load_from(&mut store)?;
// Save Burn model for PyTorch (with enum variant skipping)
let mut store = SafetensorsStore::from_file("for_pytorch.safetensors")
.with_to_adapter(BurnToPyTorchAdapter) // Convert back to PyTorch format
.skip_enum_variants(true); // Omit enum variants for PyTorch compatibility
burn_model.save_into(&mut store)?;
```
### Contiguous Layer Index Mapping
When loading PyTorch models that use `nn.Sequential` with mixed layer types (e.g., Conv2d + ReLU),
the layer indices may be non-contiguous because only some layers have parameters:
```python
# PyTorch model with non-contiguous indices
self.fc = nn.Sequential(
nn.Conv2d(...), # fc.0.weight, fc.0.bias
nn.ReLU(), # No parameters (index 1 skipped)
nn.Conv2d(...), # fc.2.weight, fc.2.bias
nn.ReLU(), # No parameters (index 3 skipped)
nn.Conv2d(...), # fc.4.weight, fc.4.bias
)
```
Burn models typically expect contiguous indices (`fc.0`, `fc.1`, `fc.2`). The `map_indices_contiguous`
feature automatically maps non-contiguous indices to contiguous ones:
```rust
// PytorchStore: map_indices_contiguous is ON by default
let mut store = PytorchStore::from_file("model.pth");
// fc.0 -> fc.0, fc.2 -> fc.1, fc.4 -> fc.2
// Disable if your model already has contiguous indices
let mut store = PytorchStore::from_file("model.pth")
.map_indices_contiguous(false);
// SafetensorsStore: map_indices_contiguous is OFF by default
let mut store = SafetensorsStore::from_file("model.safetensors")
.map_indices_contiguous(true); // Enable for PyTorch-exported safetensors
```
### Tensor Name Remapping
```rust
// Simple pattern-based remapping
let mut store = SafetensorsStore::from_file("model.safetensors")
.with_key_remapping(r"^old_model\.", "new_model.") // old_model.X -> new_model.X
.with_key_remapping(r"\.gamma$", ".weight") // X.gamma -> X.weight
.with_key_remapping(r"\.beta$", ".bias"); // X.beta -> X.bias
// Complex remapping with KeyRemapper
use burn_store::KeyRemapper;
let remapper = KeyRemapper::new()
.add_pattern(r"^transformer\.h\.(\d+)\.", "transformer.layer$1.")? // h.0 -> layer0
.add_pattern(r"^(.*?)\.attn\.", "$1.attention.")?; // attn -> attention
let mut store = SafetensorsStore::from_file("model.safetensors")
.remap(remapper);
// Combining with PyTorch loading
let mut store = PytorchStore::from_file("model.pth")
.with_key_remapping(r"^model\.", "") // Remove model. prefix
.with_key_remapping(r"norm(\d+)", "norm_$1"); // norm1 -> norm_1
```
### Memory Operations
```rust
// Burnpack: Save to memory buffer
let mut store = BurnpackStore::from_bytes(None)
.with_regex(r"^encoder\..*")
.metadata("subset", "encoder_only");
model.save_into(&mut store)?;
let bytes = store.get_bytes()?;
// Burnpack: Load from memory buffer (no-std compatible)
let mut store = BurnpackStore::from_bytes(Some(bytes))
.allow_partial(true);
let result = model.load_from(&mut store)?;
// SafeTensors: Memory operations
let mut store = SafetensorsStore::from_bytes(None)
.with_regex(r"^encoder\..*");
model.save_into(&mut store)?;
let bytes = store.get_bytes()?;
println!("Loaded {} tensors", result.applied.len());
if !result.missing.is_empty() {
println!("Missing: {:?}", result.missing);
}
```
Both BurnpackStore and SafetensorsStore support no-std environments when using byte operations
### Zero-Copy Loading
For embedded models and large model files, zero-copy loading avoids unnecessary memory allocations
by directly referencing the source data instead of copying it.
#### Embedded Models (Static Data)
```rust
use burn_store::{ModuleSnapshot, BurnpackStore};
// Embed model weights in the binary at compile time
static MODEL_DATA: &[u8] = include_bytes!("model.bpk");
// Zero-copy loading - data stays in binary's .rodata section
let mut store = BurnpackStore::from_static(MODEL_DATA);
model.load_from(&mut store)?;
```
The `from_static()` constructor automatically enables zero-copy mode. Tensor data is sliced directly
from the embedded bytes without heap allocation.
#### File-Based Zero-Copy
```rust
// Memory-mapped file with zero-copy tensor slicing
let mut store = BurnpackStore::from_file("large_model.bpk")
.zero_copy(true); // Enable zero-copy slicing
model.load_from(&mut store)?;
```
When `zero_copy(true)` is set, the memory-mapped file is wrapped in `bytes::Bytes` via
`from_owner()`, enabling O(1) slicing operations.
#### In-Memory Zero-Copy
```rust
use burn_tensor::{AllocationProperty, Bytes};
// Create shared bytes for zero-copy
let data: Vec<u8> = load_model_bytes();
let shared = bytes::Bytes::from(data);
let bytes = Bytes::from_shared(shared, AllocationProperty::Other);
// Load with zero-copy enabled
let mut store = BurnpackStore::from_bytes(Some(bytes))
.zero_copy(true);
model.load_from(&mut store)?;
```
#### When to Use Zero-Copy
| Embedded models (`include_bytes!`) | Use `from_static()` (auto-enabled) |
| Large model files | Use `from_file().zero_copy(true)` |
| Repeated loading from same bytes | Use `from_bytes().zero_copy(true)` |
| One-time load, release memory after | Use default (copy mode) |
**Note**: Zero-copy keeps the source data alive as long as any tensor references it. Use copy mode
(default) if you need to release the source file/memory immediately after loading.
### Model Surgery and Partial Operations
Burn Store enables sophisticated model surgery operations for selectively loading, saving, and
transferring parts of models.
#### Direct Model-to-Model Transfer
```rust
use burn_store::{ModuleSnapshot, PathFilter};
// Direct transfer - all compatible tensors
let snapshots = model1.collect(None, None, false);
let result = model2.apply(snapshots, None, None, false);
// Selective transfer with filtering
let filter = PathFilter::new().with_regex(r"^encoder\..*");
let snapshots = model1.collect(Some(filter.clone()), None, false);
let result = model2.apply(snapshots, Some(filter), None, false);
// Transfer with path transformation
let mut snapshots = model1.collect(None, None, false);
for snapshot in &mut snapshots {
snapshot.full_path = snapshot.full_path.replace("encoder.", "transformer.encoder.");
}
model2.apply(snapshots, None, None, false);
```
#### Partial Loading and Exports
```rust
// Export only specific layers
let mut store = SafetensorsStore::from_file("encoder_only.safetensors")
.with_regex(r"^encoder\..*");
model.save_into(&mut store)?;
// Load with missing tensors allowed
let mut store = SafetensorsStore::from_file("pretrained.safetensors")
.allow_partial(true);
let result = model.load_from(&mut store)?;
println!("Loaded: {}, Missing: {:?}", result.applied.len(), result.missing);
```
#### Merging Multiple Models
```rust
// Merge weights from different sources
let mut merged = Vec::new();
merged.extend(base_model.collect(None, None, false));
// Add encoder from specialized model
let encoder_filter = PathFilter::new().with_regex(r"^encoder\..*");
merged.extend(specialized_model.collect(Some(encoder_filter), None, false));
// Apply merged weights
target_model.apply(merged, None, None, false);
// Alternative: Sequential loading from files
let mut base_store = SafetensorsStore::from_file("base.safetensors");
model.load_from(&mut base_store)?;
let mut encoder_store = SafetensorsStore::from_file("encoder.safetensors")
.with_regex(r"^encoder\..*")
.allow_partial(true);
model.load_from(&mut encoder_store)?; // Overlays encoder weights
```
### Complete Example: Migrating PyTorch Models
```rust
use burn_store::{ModuleSnapshot, PytorchStore};
// Load directly from PyTorch .pth file (automatic PyTorchToBurnAdapter)
// Note: skip_enum_variants defaults to true for PytorchStore
let mut store = PytorchStore::from_file("pytorch_transformer.pth")
// Access the state dict
.with_top_level_key("state_dict")
// Only load transformer layers
.with_regex(r"^transformer\..*")
// Rename layer structure to match Burn model
.with_key_remapping(r"^transformer\.h\.(\d+)\.", "transformer.layer$1.")
// Rename attention layers
.with_key_remapping(r"\.attn\.", ".attention.")
// Handle missing tensors gracefully
.allow_partial(true);
let mut model = TransformerModel::new(&device);
let result = model.load_from(&mut store)?;
println!("Successfully migrated {} tensors", result.applied.len());
if !result.errors.is_empty() {
println!("Errors: {:?}", result.errors);
}
// Save the migrated model in SafeTensors format
let mut save_store = SafetensorsStore::from_file("migrated_model.safetensors")
.metadata("source", "pytorch")
.metadata("converted_by", "burn-store");
model.save_into(&mut save_store)?;
```
## Advanced Usage
### Direct Tensor Access
All stores provide methods to directly access tensor snapshots without loading into a model. This is
useful for inspection, debugging, selective processing, or building custom loading pipelines.
```rust
use burn_store::{ModuleStore, BurnpackStore, SafetensorsStore, PytorchStore};
// Works with any store type
let mut store = BurnpackStore::from_file("model.bpk");
// let mut store = SafetensorsStore::from_file("model.safetensors");
// let mut store = PytorchStore::from_file("model.pth");
// List all tensor names (ordered)
let names = store.keys()?;
println!("Model contains {} tensors:", names.len());
for name in &names {
println!(" - {}", name);
}
// Get all tensors as a BTreeMap (cached for repeated access)
let snapshots = store.get_all_snapshots()?;
for (name, snapshot) in snapshots {
println!("{}: {:?} {:?}", name, snapshot.shape, snapshot.dtype);
}
// Get a specific tensor by name
if let Some(snapshot) = store.get_snapshot("encoder.layer0.weight")? {
// Lazy loading - data is only fetched when to_data() is called
let data = snapshot.to_data()?;
println!("Shape: {:?}, DType: {:?}", data.shape, data.dtype);
}
```
#### Use Cases
- **Model Inspection**: Examine tensor shapes, dtypes, and names without full model instantiation
- **Selective Loading**: Build custom pipelines that only load specific tensors
- **Debugging**: Verify tensor values and compare across different model files
- **Format Conversion**: Read tensors from one format and write to another
### Custom Filtering with Predicates
```rust
// Custom filter function
let mut store = SafetensorsStore::from_file("model.safetensors")
.with_predicate(|path, _container| {
// Only load tensors with specific characteristics
path.contains("weight") && !path.contains("bias")
});
```
### Working with Containers
```rust
// Filter based on container types (Linear, Conv2d, etc.)
let mut store = SafetensorsStore::from_file("model.safetensors")
.with_predicate(|_path, container| {
// Only load Linear layer parameters
container.split('.').last() == Some("Linear")
});
```
### Handling Load Results
```rust
let result = model.load_from(&mut store)?;
// Detailed result information
println!("Applied: {} tensors", result.applied.len());
println!("Skipped: {} tensors", result.skipped.len());
println!("Missing: {:?}", result.missing);
println!("Unused: {:?}", result.unused);
if !result.errors.is_empty() {
for (path, error) in &result.errors {
eprintln!("Error loading {}: {}", path, error);
}
}
```
## Benchmarks
### Loading Benchmarks
```bash
# Generate model files first (one-time setup)
cd crates/burn-store
uv run benches/generate_unified_models.py
# Run unified loading benchmark with default backend (NdArray CPU)
cargo bench --bench unified_loading
# Run with specific backend
cargo bench --bench unified_loading --features metal # Apple GPU
cargo bench --bench unified_loading --features wgpu # WebGPU
cargo bench --bench unified_loading --features cuda # NVIDIA GPU
cargo bench --bench unified_loading --features candle # Candle backend
cargo bench --bench unified_loading --features tch # LibTorch
# Run with multiple backends
cargo bench --bench unified_loading --features wgpu,tch
```
### Saving Benchmarks
Compares 3 saving methods: BurnpackStore, NamedMpkFileRecorder, and SafetensorsStore.
```bash
# Run unified saving benchmark with default backend (NdArray CPU)
cargo bench --bench unified_saving
# Run with specific backend
cargo bench --bench unified_saving --features metal # Apple GPU
cargo bench --bench unified_saving --features wgpu # WebGPU
cargo bench --bench unified_saving --features cuda # NVIDIA GPU
cargo bench --bench unified_saving --features candle # Candle backend
cargo bench --bench unified_saving --features tch # LibTorch
# Run with multiple backends
cargo bench --bench unified_saving --features wgpu,tch
```
## API Overview
### Builder Methods
The stores provide a fluent API for configuration:
#### Filtering
- `with_regex(pattern)` - Filter by regex pattern
- `with_full_path(path)` - Include specific tensor
- `with_full_paths(paths)` - Include multiple specific tensors
- `with_predicate(fn)` - Custom filter logic
- `match_all()` - Include all tensors (no filtering)
#### Remapping
- `with_key_remapping(from, to)` - Regex-based tensor renaming
- `remap(KeyRemapper)` - Complex remapping rules
#### Adapters
- `with_from_adapter(adapter)` - Loading transformations
- `with_to_adapter(adapter)` - Saving transformations
#### Configuration
- `metadata(key, value)` - Add custom metadata (Burnpack and SafeTensors)
- `allow_partial(bool)` - Continue on missing tensors
- `validate(bool)` - Toggle validation
- `skip_enum_variants(bool)` - Skip enum variant names in paths for PyTorch compatibility
- `map_indices_contiguous(bool)` - Map non-contiguous layer indices to contiguous (default: `true`
for PyTorch, `false` for SafeTensors)
- `with_top_level_key(key)` - Access nested dict in PyTorch files
- `overwrite(bool)` - Allow overwriting existing files (Burnpack)
- `zero_copy(bool)` - Enable zero-copy tensor slicing (Burnpack)
#### Direct Tensor Access
- `keys()` - Get ordered list of all tensor names
- `get_all_snapshots()` - Get all tensors as a BTreeMap (cached)
- `get_snapshot(name)` - Get a specific tensor by name
### Inspecting Burnpack Files
Generate and examine a sample file:
```bash
cargo run --example burnpack_inspect sample.bpk
The example creates a sample model and outputs inspection commands for examining the binary format.
## License
This project is dual-licensed under MIT and Apache-2.0.