# Migration Guide: burn-import to burn-store
This guide helps you migrate from the deprecated `burn-import` recorders (`PyTorchFileRecorder`,
`SafetensorsFileRecorder`) to the new `burn-store` API (`PytorchStore`, `SafetensorsStore`).
## Overview
The new `burn-store` API provides:
- **Simpler API**: Load directly into models instead of records
- **Fluent builder pattern**: Chain configuration methods
- **Better error handling**: Detailed load results with applied/missing/errors info
- **Bidirectional support**: Both load and save operations
- **More features**: Filtering, partial loading, metadata, zero-copy loading
## Quick Migration
### PyTorch Files (.pt/.pth)
**Before (burn-import):**
```rust
use burn::record::{FullPrecisionSettings, Recorder};
use burn_import::pytorch::{LoadArgs, PyTorchFileRecorder};
// Load into a record, then create model from record
let record: ModelRecord<B> = PyTorchFileRecorder::<FullPrecisionSettings>::default()
.load("model.pt".into(), &device)
.expect("Failed to load");
let model = Model::init(&device).load_record(record);
```
**After (burn-store):**
```rust
use burn_store::{ModuleSnapshot, PytorchStore};
// Initialize model, then load weights directly
let mut model = Model::init(&device);
let mut store = PytorchStore::from_file("model.pt");
model.load_from(&mut store).expect("Failed to load");
```
### SafeTensors Files (.safetensors)
**Before (burn-import):**
```rust
use burn::record::{FullPrecisionSettings, Recorder};
use burn_import::safetensors::{AdapterType, LoadArgs, SafetensorsFileRecorder};
let record: ModelRecord<B> = SafetensorsFileRecorder::<FullPrecisionSettings>::default()
.load("model.safetensors".into(), &device)
.expect("Failed to load");
let model = Model::init(&device).load_record(record);
```
**After (burn-store):**
```rust
use burn_store::{ModuleSnapshot, PyTorchToBurnAdapter, SafetensorsStore};
let mut model = Model::init(&device);
// For SafeTensors exported from PyTorch, use the adapter
let mut store = SafetensorsStore::from_file("model.safetensors")
.with_from_adapter(PyTorchToBurnAdapter);
model.load_from(&mut store).expect("Failed to load");
// For native Burn SafeTensors, no adapter needed
let mut store = SafetensorsStore::from_file("model.safetensors");
model.load_from(&mut store).expect("Failed to load");
```
## API Mapping
### PyTorchFileRecorder Options
| `LoadArgs::new(path)` | `PytorchStore::from_file(path)` |
| `.with_key_remap(pattern, replacement)` | `.with_key_remapping(pattern, replacement)` |
| `.with_top_level_key(key)` | `.with_top_level_key(key)` |
| `.with_debug_print()` | _(use tracing/logging instead)_ |
| `PyTorchFileRecorder::<FullPrecisionSettings>` | _(precision handled automatically)_ |
### SafetensorsFileRecorder Options
| `LoadArgs::new(path)` | `SafetensorsStore::from_file(path)` |
| `.with_key_remap(pattern, replacement)` | `.with_key_remapping(pattern, replacement)` |
| `.with_adapter_type(AdapterType::PyTorch)` | `.with_from_adapter(PyTorchToBurnAdapter)` |
| `.with_adapter_type(AdapterType::NoAdapter)` | _(default, no adapter)_ |
| `.with_debug_print()` | _(use tracing/logging instead)_ |
| `SafetensorsFileRecorder::<FullPrecisionSettings>` | _(precision handled automatically)_ |
## Detailed Examples
### Key Remapping
**Before:**
```rust
let args = LoadArgs::new("model.pt".into())
.with_key_remap("conv\\.(.*)", "$1")
.with_key_remap("^old_prefix\\.", "new_prefix.");
let record: ModelRecord<B> = PyTorchFileRecorder::<FullPrecisionSettings>::default()
.load(args, &device)?;
```
**After:**
```rust
let mut store = PytorchStore::from_file("model.pt")
.with_key_remapping("conv\\.(.*)", "$1")
.with_key_remapping("^old_prefix\\.", "new_prefix.");
model.load_from(&mut store)?;
```
### Top-Level Key Access
**Before:**
```rust
let args = LoadArgs::new("checkpoint.pt".into())
.with_top_level_key("state_dict");
let record: ModelRecord<B> = PyTorchFileRecorder::<FullPrecisionSettings>::default()
.load(args, &device)?;
```
**After:**
```rust
let mut store = PytorchStore::from_file("checkpoint.pt")
.with_top_level_key("state_dict");
model.load_from(&mut store)?;
```
### PyTorch Adapter for SafeTensors
**Before:**
```rust
use burn_import::safetensors::{AdapterType, LoadArgs};
let args = LoadArgs::new("pytorch_model.safetensors".into())
.with_adapter_type(AdapterType::PyTorch);
let record: ModelRecord<B> = SafetensorsFileRecorder::<FullPrecisionSettings>::default()
.load(args, &device)?;
```
**After:**
```rust
use burn_store::{PyTorchToBurnAdapter, SafetensorsStore};
let mut store = SafetensorsStore::from_file("pytorch_model.safetensors")
.with_from_adapter(PyTorchToBurnAdapter);
model.load_from(&mut store)?;
```
## New Features in burn-store
### Partial Loading
Handle missing tensors gracefully:
```rust
let mut store = PytorchStore::from_file("model.pt")
.allow_partial(true);
let result = model.load_from(&mut store)?;
println!("Loaded: {:?}", result.applied);
println!("Missing: {:?}", result.missing);
```
### Filtering
Load only specific tensors:
```rust
let mut store = SafetensorsStore::from_file("model.safetensors")
.with_regex(r"^encoder\..*") // Only encoder layers
.allow_partial(true);
model.load_from(&mut store)?;
```
### Saving Models
Save models (not supported by old recorders):
```rust
// Save to SafeTensors
let mut store = SafetensorsStore::from_file("output.safetensors")
.metadata("version", "1.0");
model.save_into(&mut store)?;
// Save to Burnpack (native format)
let mut store = BurnpackStore::from_file("output.bpk");
model.save_into(&mut store)?;
```
### Load Results
Get detailed information about loading:
```rust
let result = model.load_from(&mut store)?;
// Print the full result for debugging - shows applied, skipped, missing, and errors
println!("{}", result);
// Or access individual fields
println!("Applied: {} tensors", result.applied.len());
println!("Skipped: {} tensors", result.skipped.len());
println!("Missing: {:?}", result.missing);
println!("Errors: {:?}", result.errors);
// Check if fully successful
if result.is_success() {
println!("All tensors loaded successfully");
}
```
The `LoadResult` implements `Display`, so printing it shows a formatted summary with suggestions for
common issues (e.g., using `allow_partial(true)` for missing tensors).
## Updating Cargo.toml
**Before:**
```toml
[dependencies]
burn-import = { version = "0.x", features = ["pytorch", "safetensors"] }
```
**After:**
```toml
[dependencies]
burn-store = { version = "0.x", features = ["pytorch", "safetensors"] }
```
## Common Migration Issues
### 1. Model vs Record
The new API loads directly into models, not records. Update your model initialization:
```rust
// Before: Create record, then model from record
let record = recorder.load(...)?;
let model = Model::init(&device).load_record(record);
// After: Create model, then load into it
let mut model = Model::init(&device);
model.load_from(&mut store)?;
```
### 2. Inference Functions
If you had functions that took `ModelRecord`, update them to take `Model`:
```rust
// Before
fn infer(record: ModelRecord<B>) {
let model = Model::init(&device).load_record(record);
// ...
}
// After
fn infer(model: Model<B>) {
// Model already has weights loaded
// ...
}
```
### 3. Precision Settings
The old API required explicit precision settings. The new API handles this automatically:
```rust
// Before: Had to specify FullPrecisionSettings or HalfPrecisionSettings
PyTorchFileRecorder::<FullPrecisionSettings>::default()
// After: Precision handled automatically based on tensor dtype
PytorchStore::from_file("model.pt")
```
### 4. Error Handling
The new API provides richer error information:
```rust
// Before: Simple Result
let record = recorder.load(args, &device)?;
// After: LoadResult with detailed info
let result = model.load_from(&mut store)?;
// Print the result to see a helpful summary with suggestions
println!("{}", result);
// Or handle specific issues programmatically
if !result.errors.is_empty() {
for (path, error) in &result.errors {
eprintln!("Error loading {}: {}", path, error);
}
}
```
## See Also
- [burn-store README](README.md) - Full documentation
- [import-model-weights example](../../examples/import-model-weights/) - Working example