Rust in Everything
基于 burn 0.22.1

保存与加载

主线位置：部署主线的前两步，保存和加载。训练好的参数写进文件，另一个进程先建出同样结构的模型，再把文件里的参数套进去。
版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）
flowchart TB
  model["模型对象"] -->|"按字段路径收集"| table["名字到张量的表"]
  table --> bpk[".bpk burnpack"]
  table --> st[".safetensors"]
  fresh["新建同结构模型"] --> applyn["按名字套回"]
  bpk --> applyn
  st --> applyn
  applyn --> check{"缺 多 形状不符"}
  check -->|"都没有"| ok["可以推理"]
  check -->|"有"| err["报错"]

模型用主线的两层 MLP（4→8→3，ReLU，权重按固定公式生成，见怎么读这一卷）。实验代码：PyTorch 侧 save_load.py，burn 侧 crate store-io 的 bin saveload（default-features = false，feature std、flex、autodiff、optim、store、safetensors、pytorch，debug）。平台 M2 Max、macOS 27、rustc 1.99.0。

存的是一张「名字 → 张量」的表

PyTorch 的 state_dict() 把参数按属性路径列出来：

sd = model.state_dict()
for k, v in sd.items():
  print(f"{k}: {list(v.shape)} {v.dtype}")

fc1.weight: [8, 4] torch.float32
fc1.bias: [8] torch.float32
fc2.weight: [3, 8] torch.float32
fc2.bias: [3] torch.float32

burn 的 #[derive(Module)] 生成的遍历按字段路径收集参数（见定义模型），名字和 PyTorch 一样是 fc1.weight，形状按 burn 自己的约定。用 burn-store 把主线模型原样（不加适配器）写成 safetensors，再在 Python 里列出来：

burn_a1_raw metadata: {'producer': 'burn', 'version': '0.22.1', 'format': 'safetensors'}
burn_a1_raw {'fc1.bias': [8], 'fc1.weight': [4, 8], 'fc2.bias': [3], 'fc2.weight': [8, 3]}

名字相同，两个权重是转置的：burn 的 Linear.weight 是 [d_input, d_output]。这是两边存取时唯一要换的地方，换法在后面两节。

两边都不把模型结构存进文件。加载之前，先在代码里建出结构相同的模型：PyTorch 是 Mlp()，burn 是 Mlp::new(&device)，参数先是随机初始化的值，加载时被文件里的值替换。

📝NOTE
没有蠢问题

问： PyTorch 还能 torch.save(model) 把整个模型存下来，burn 有没有对应的写法？

答： 没有。torch.save(model) 存的是 pickle：参数之外，还有模型类的模块名和类名，读回来时 Python 按这些名字找到类再重建对象，所以读的一方仍然要能 import 到同样的类定义：在一个找不到定义 Mlp 的模块 a1 的进程里，torch.load("a1_full.pt", weights_only=False) 报 ModuleNotFoundError: No module named 'a1'。burn 的模型是 Rust 类型，结构和 forward 编译进程序，文件里只放参数。要交付一个不带模型源码的计算图，PyTorch 有 TorchScript 和 torch.export，burn 0.22.1 的情况见部署形态。

存一次、读回来

PyTorch：

model = a1_model()
torch.save(model.state_dict(), f"{out}/rt.pt")
fresh = Mlp()
print(fresh.load_state_dict(torch.load(f"{out}/rt.pt")))
with torch.no_grad():
  print("equal:", torch.equal(fresh(x), model(x)))

<All keys matched successfully>
equal: True

burn：

let model = a1_model(&device);
model.clone().save_file(format!("{dir}/a1.bpk"))?;   // "save" 模式
// ...
let fresh = Mlp::new(&device);                         // "load" 模式，另一个进程
let model = fresh.try_load_file(format!("{dir}/a1.bpk"))?;

（实验代码里的 ? 是 match 加打印错误，这里简写。）

loaded params bit-equal to saved: true
logits bit-equal to the model before saving: true

读回来的参数和保存前逐位相同，前向结果也逐位相同。两边概念一样，写法上有两处不同：

  • save_file 和 load_file 按值拿走模型：load_file(self, path) -> Self 返回一个装好参数的新模型，而不是像 load_state_dict 那样原地改。保存后还要用模型就先 clone()，它只复制参数的引用（见运算消费张量）。
  • load_file 出错就 panic，try_load_file 返回 Result；torch.load 出错抛异常。

save_file 写的是 burn 自己的 burnpack 格式（.bpk），在 burn-core 里，不用开任何额外 feature。它除了张量还存每个参数的 ParamId，加载后恢复，所以按 ParamId 存的优化器状态读回来还能对上参数；格式和这一点的细节见 burn 卷的存储与加载。PyTorch 的 .pt 文件 burn 只能读、不能写，下一节和导入 PyTorch 权重讲。

严格加载：缺的、多的、形状不符

load_state_dict 默认 strict=True。把文件里的 fc1.weight 改名为 encoder.weight，再单独试一次把 fc1.weight 转置：

bad = dict(sd)
bad["encoder.weight"] = bad.pop("fc1.weight")
Mlp().load_state_dict(bad)                   # 抛 RuntimeError
print(Mlp().load_state_dict(bad, strict=False))
tr = dict(sd)
tr["fc1.weight"] = tr["fc1.weight"].T.contiguous()
Mlp().load_state_dict(tr)                    # 抛 RuntimeError

RuntimeError: Error(s) in loading state_dict for Mlp:
	Missing key(s) in state_dict: "fc1.weight". 
	Unexpected key(s) in state_dict: "encoder.weight". 
_IncompatibleKeys(missing_keys=['fc1.weight'], unexpected_keys=['encoder.weight'])
RuntimeError: Error(s) in loading state_dict for Mlp:
	size mismatch for fc1.weight: copying a param with shape torch.Size([4, 8]) from checkpoint, the shape in current model is torch.Size([8, 4]).

burn 也查这三类问题，但两条加载路径的默认不一样：

情况PyTorch load_state_dictburn load_fileburn-store load_from
模型要的名字文件里没有报错（strict=False 时放过）报错报错（allow_partial(true) 时放过）
文件里有、模型不要报错（strict=False 时放过）报错不报错，只列在摘要里
形状不符报错，strict=False 也报报错报错（ShapeMismatch）

改了字段名之后 burn 的报错原文和改法见 burn 卷的 Record validation error: Missing tensors，ShapeMismatch 见导入 PyTorch 权重。strict=False 返回的 _IncompatibleKeys 对应 burn-store 的 ApplyResult：load_from 成功时返回它，打印出来是一张列出缺失、多余、错误的表。

safetensors：burn 写、PyTorch 读

两边都能读写 safetensors，它是交换参数的格式。burn 侧要开 store 和 safetensors 两个 feature，用 burn-store 的 SafetensorsStore。写给 PyTorch 读时挂一个 BurnToPyTorchAdapter，把 Linear 的权重转回 [out, in]（with_to_adapter 管保存方向）：

use burn::store::{BurnToPyTorchAdapter, ModuleSnapshot, SafetensorsStore};

let mut st = SafetensorsStore::from_file(format!("{dir}/burn_a1.safetensors"))
  .with_to_adapter(BurnToPyTorchAdapter);
println!("save_into (BurnToPyTorchAdapter): {:?}", model.save_into(&mut st).map(|_| "Ok"));
// 不加适配器：权重按 burn 的 [in, out] 写出
let mut st = SafetensorsStore::from_file(format!("{dir}/burn_a1_raw.safetensors"));
println!("save_into (no adapter): {:?}", model.save_into(&mut st).map(|_| "Ok"));

save_into (BurnToPyTorchAdapter): Ok("Ok")
save_into (no adapter): Ok("Ok")

PyTorch 侧用 safetensors.torch.load_file 读，再 load_state_dict：

for name in ["burn_a1", "burn_a1_raw"]:
  path = f"{out}/{name}.safetensors"
  with safe_open(path, "pt") as f:
    print(name, "metadata:", f.metadata())
    print(name, {k: list(f.get_slice(k).get_shape()) for k in f.keys()})
  m = Mlp()
  try:
    print(name, m.load_state_dict(load_file(path)))
  except RuntimeError as e:
    print(name, "RuntimeError:", e)

burn_a1 metadata: {'format': 'safetensors', 'version': '0.22.1', 'producer': 'burn'}
burn_a1 {'fc1.bias': [8], 'fc1.weight': [8, 4], 'fc2.bias': [3], 'fc2.weight': [3, 8]}
burn_a1 <All keys matched successfully>
burn_a1_raw metadata: {'producer': 'burn', 'version': '0.22.1', 'format': 'safetensors'}
burn_a1_raw {'fc1.bias': [8], 'fc1.weight': [4, 8], 'fc2.bias': [3], 'fc2.weight': [8, 3]}
burn_a1_raw RuntimeError: Error(s) in loading state_dict for Mlp:
	size mismatch for fc1.weight: copying a param with shape torch.Size([4, 8]) from checkpoint, the shape in current model is torch.Size([8, 4]).
	size mismatch for fc2.weight: copying a param with shape torch.Size([8, 3]) from checkpoint, the shape in current model is torch.Size([3, 8]).

加了适配器的文件 PyTorch 直接能读。另外把 burn 模型的参数按 [out, in] 打印成 JSON，和 PyTorch 从 burn_a1.safetensors 读到的张量逐个 torch.equal，四个都是 True：写出去的是同一组 f32，没有任何舍入。不加适配器的文件，PyTorch 一侧报 size mismatch。

元数据里的 format、producer、version 三项是 default_metadata 自动加的。

两个方向不对称。burn 写 safetensors 默认按 burn 的布局写，要给 PyTorch 读得自己挂 BurnToPyTorchAdapter；burn 读 PyTorch 的 safetensors 也要自己挂 PyTorchToBurnAdapter；只有读 .pt 的 PytorchStore 自带适配器。反过来写 .pt 不行：

let mut st = burn::store::PytorchStore::from_file(format!("{dir}/burn_a1.pt"));
match model.save_into(&mut st) { /* ... */ }

PytorchStore save_into: Err Saving to PyTorch format is not yet supported. Use other formats for saving.

错误来自 PytorchStore 的 collect_from，它不看模型，直接返回这个错误。

同一个路径存第二次

torch.save 存到已有的路径，直接覆盖：

second torch.save: ok

burn 的两条路不一样：

let mut st = SafetensorsStore::from_file(format!("{dir}/burn_a1.safetensors"))
  .with_to_adapter(BurnToPyTorchAdapter);
match model.save_into(&mut st) { /* second save_into */ }
let mut st = SafetensorsStore::from_file(format!("{dir}/burn_a1.safetensors"))
  .with_to_adapter(BurnToPyTorchAdapter)
  .overwrite(true);
match model.save_into(&mut st) { /* overwrite(true) */ }
match model.clone().save_file(format!("{dir}/a1.bpk")) { /* second save_file */ }

second save_into: Err File already exists: ../store-out/burn_a1.safetensors. To replace it, use .overwrite(true), or AtomicFile::commit instead of commit_new.
overwrite(true): Ok
second save_file: Ok

burn-store 的 SafetensorsStore 和 BurnpackStore 默认 overwrite: false，文件已经存在就返回错误，要显式 overwrite(true)。Module::save_file 走 burn-core 的另一条路，和 torch.save 一样直接覆盖。训练循环里每个 epoch 存到同一个文件名的写法，换到 burn-store 上第二个 epoch 就会报错。

📝NOTE
冷知识：「不许覆盖」以前只在开始时查一次

#5781「fix(store): enforce overwrite(false) at publish time, not just before the save」（首次发布于 v0.22.0）的提交说明写了之前的问题：「The overwrite guard checked for the destination before the save and then renamed over it, so a file created during the save was silently replaced.」也就是检查和写入之间有一个窗口，保存期间别的进程写出的文件会被悄悄换掉。改法是在 burn-pack 里加一个不覆盖的发布方式：AtomicFile::commit_new 用硬链接把写好的临时文件放到目标路径，目标已存在时链接失败，返回新的 Error::AlreadyExists。开头的检查保留，只用来在读回任何张量之前就早点失败。上面报错里的「or AtomicFile::commit instead of commit_new」就来自同一个 PR，它让这条提示也适用于直接用 AtomicFile 的调用者。先写临时文件再改名这一步来自更早的 #5489（v0.22.0-pre.4）：之前 serialize_to_file 一开始就截断目标文件，张量读回中途失败时，原来那份好的 checkpoint 也没了。

小结

  • 两边存的都是「参数路径 → 张量」的表，名字都是 fc1.weight 这样的路径；结构不进文件，加载前先建模型。
  • torch.save(state_dict) / load_state_dict 对 save_file / load_file；burn 的加载按值拿走模型、返回新模型，读回的参数和前向结果逐位相同。
  • 严格加载：PyTorch 和 load_file 对缺失、多余、形状不符都报错；burn-store 的 load_from 对多余的张量只在摘要里列出。
  • 和 PyTorch 交换用 safetensors：burn 写时挂 BurnToPyTorchAdapter，PyTorch 直接 load_state_dict，数值逐位相同；不挂就是 size mismatch。burn 不能写 .pt。
  • burn-store 默认不覆盖已有文件，要 overwrite(true)；save_file 和 torch.save 一样直接覆盖。

💬 关联讨论

发起讨论
加载中...