导入 PyTorch 权重
主线位置：部署主线的「加载」一步，参数来自 PyTorch。模型在 PyTorch 里训练好，burn 侧用 Rust 重写同样结构的模型，把权重文件读进来推理。
版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）
flowchart TB
pyf["PyTorch 存的文件"] --> kind{"哪种文件"}
kind -->|".safetensors"| sts["SafetensorsStore"]
kind -->|".pt"| pts["PytorchStore"]
sts -->|"自己挂"| adapt["PyTorchToBurnAdapter"]
pts -->|"自带"| adapt
adapt --> names{"键名对得上"}
names -->|"Sequential 编号"| remap["改名或重编号"]
names -->|"对上"| dtype{"dtype 一样"}
remap --> dtype
dtype -->|"F64 文件"| cast["FloatCastAdapter"]
dtype -->|"F32"| fwd["forward 对照"]
cast --> fwd
实验代码：PyTorch 侧 export.py（把主线模型存成几种文件，并打印前向结果），burn 侧 crate store-io 的 bin import（参数选情形；default-features = false，feature std、flex、autodiff、optim、store、safetensors、pytorch，debug）。平台 M2 Max、macOS 27、rustc 1.99.0。存取的基本写法见保存与加载。
PyTorch 侧：存成 safetensors
主线 MLP，权重按固定公式生成，eval()，存 state_dict：
model = a1_model()
sd = model.state_dict()
save_file(sd, f"{out}/a1.safetensors") # safetensors.torch.save_file
torch.save(sd, f"{out}/a1_state.pt")
torch.save(sd, f"{out}/a1_state_legacy.pt", _use_new_zipfile_serialization=False)
torch.save(model, f"{out}/a1_full.pt")
fc1.weight: [8, 4] torch.float32
fc1.bias: [8] torch.float32
fc2.weight: [3, 8] torch.float32
fc2.bias: [3] torch.float32
burn 侧的模型和主线一样，Mlp { fc1: Linear, fc2: Linear, relu: Relu }，字段名和 PyTorch 的属性名一致，键名就对得上。
直接读：ShapeMismatch
let mut model = Mlp::new(&device);
let mut st = SafetensorsStore::from_file(format!("{dir}/a1.safetensors"));
match model.load_from(&mut st) {
Ok(r) => println!("Ok\n{r}"),
Err(e) => println!("Err\n{e}"),
}
Err
Validation failed: Import errors: [ShapeMismatch { path: "fc1.weight", expected: Shape { dims: [4, 8] }, found: Shape { dims: [8, 4] } }, ShapeMismatch { path: "fc2.weight", expected: Shape { dims: [8, 3] }, found: Shape { dims: [3, 8] } }]
PyTorch 的 nn.Linear 算 x @ W.T + b，weight 是 [out, in]；burn 的 Linear 算 x @ W + b，weight 是 [in, out]。burn-store 只比形状，不按名字猜哪个张量要转置，所以两个权重各报一条，偏置是一维的，对上了。这个错误的来龙去脉见 burn 卷的 ShapeMismatch expected [4, 3] found [3, 4]。
加上 PyTorchToBurnAdapter：前向对得上
let mut st = SafetensorsStore::from_file(format!("{dir}/a1.safetensors"))
.with_from_adapter(PyTorchToBurnAdapter);
match model.load_from(&mut st) { /* 打印结果 */ }
println!("fc1.weight dims {:?}", model.fc1.weight.val().dims());
Ok
┌─ Tensor Loading Summary ─────────────────────────
│
│ ✓ Successfully applied: 4 tensors
│ ⊘ Skipped (filtered): 0 tensors
│ ✗ Missing in source: 0 tensors
│ ? Unused in target: 0 tensors
│ ! Errors: 0 errors
│
└───────────────────────────────────────────────────
fc1.weight dims [4, 8]
PyTorchToBurnAdapter 按参数所在的模块类型决定做什么：属于 Linear 的 weight 转置，归一化层的 weight / bias 改名为 gamma / beta，别的不动。
再用同样的 4 个样本跑一次前向。输入按 Python 的算法生成：0.1 * k - 0.5 先在 f64 里算，再舍入成 f32。
let xs: Vec<f32> = (0..4 * IN).map(|k| (0.1 * k as f64 - 0.5) as f32).collect();
let x = Tensor::<2>::from_data(TensorData::new(xs, [4, IN]), &device);
let logits = model.forward(x);
{"logits": [[-0.09422215819358826, -0.12721742689609528, 0.08542244881391525], [0.0261688269674778, -0.12934733927249908, -0.0343487448990345], [0.056371480226516724, -0.20086252689361572, -0.04374047741293907], [0.11677984893321991, -0.35583826899528503, -0.059050895273685455]]}
{"logits":[[-0.09422215819358826,-0.12721742689609528,0.08542244881391525],[0.02616882510483265,-0.12934733927249908,-0.0343487411737442],[0.056371480226516724,-0.20086251199245453,-0.043740492314100266],[0.11677984893321991,-0.35583826899528503,-0.059050895273685455]]}
logits n= 12 max|diff|=1.49e-08
worst=1.49e-08 tol=1e-06 OK
上面是 PyTorch，下面是 burn。参数是同一组 f32（从 PyTorch 的文件读来），输入也逐位相同。12 个数里 8 个逐位相同，3 个差 1 个 ulp，1 个（-0.0437）差 4 个 ulp，也就是 1.49e-8，这是最大误差，容差 1e-6。不是 0 的原因是两边 matmul 内部的求和顺序不同，见怎么读这一卷。
输入要按 Python 的方式生成，这一点容易漏：写成 0.1 * k as f32 - 0.5（全程 f32）时，16 个输入里有 7 个和 PyTorch 差一个 ulp，最大 5.96e-8。这一次碰巧没有改变最大误差（仍是 1.49e-8），但不同的 logits 从 4 个变成 6 个；对照时先确认输入逐位相同，误差才只来自计算本身。
方阵：加载成功，结果是错的
上面的 MLP 没有方阵权重，形状不符一定会报错。in == out 的 Linear 不会。PyTorch 里建一个 nn.Linear(4, 4)，权重是 arange(16).reshape(4, 4)、偏置 0，存成 safetensors：
sq = nn.Linear(4, 4)
with torch.no_grad():
sq.weight.copy_(torch.arange(16.0).reshape(4, 4))
sq.bias.zero_()
save_file(sq.state_dict(), f"{out}/sq.safetensors")
print("sq forward([1,0,0,0]) =", sq(torch.tensor([[1.0, 0.0, 0.0, 0.0]])).flatten().tolist())
sq forward([1,0,0,0]) = [0.0, 4.0, 8.0, 12.0]
burn 侧分别不加和加适配器读进 LinearConfig::new(4, 4)：
no adapter: ok = true, forward([1,0,0,0]) = [0.0, 1.0, 2.0, 3.0]
adapter: ok = true, forward([1,0,0,0]) = [0.0, 4.0, 8.0, 12.0]
不加适配器也「加载成功」，算出来的是权重的第 0 行，PyTorch 算的是第 0 列。真实模型里注意力的 q_proj、o_proj 常常是方阵，所以导入 PyTorch 权重时不管报不报错都挂上适配器，加载后用一个输入和 PyTorch 的输出逐个数比一次。
.pt 文件：PytorchStore
PyTorch 自己的 .pt 文件用 PytorchStore 读。它的 apply_to 总是用 PyTorchToBurnAdapter，不用自己挂：
let mut model = Mlp::new(&device);
let mut st = PytorchStore::from_file(format!("{dir}/{file}"));
match model.load_from(&mut st) { /* ... */ }
上面存的三种 .pt 都能读，Successfully applied: 4 tensors，前向和 PyTorch 比：
| 文件 | 怎么存的 | 最大误差 |
|---|---|---|
a1_state.pt | torch.save(state_dict)，默认的 ZIP 格式 | 1.49e-8 |
a1_state_legacy.pt | _use_new_zipfile_serialization=False，旧格式 | 1.49e-8 |
a1_full.pt | torch.save(model)，整个模型 | 1.49e-8 |
（这三行用的是 f32 全程算出的输入，最大误差同样是 1.49e-8。）
第三个文件在 PyTorch 2.15.0 里反而不能直接 torch.load：
（原文里有一段加粗用的终端控制符，这里去掉了。）
UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
(1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
...
PytorchStore 读它不需要 Python，也不执行文件里的东西。它底下的 pytorch-reader 自己解析 pickle，reduce 只解释 PyTorch 重建张量的那几个调用和 OrderedDict、set 的构造，别的调用当作不透明对象留下；整个模型的 pickle 按它的 state_dict() 读。读 .pt 的更多细节（容器格式、懒读取）见 burn 卷的存储与加载。那一页还记录了一个上游测试数据里的 ZIP 格式 .pt 用 zip 9.1.3 打不开；本页的 crate 解析到的也是 zip 9.1.3，PyTorch 2.15.0 新写出的两个 ZIP 格式文件都能读。
键名：nn.Sequential 的编号
PyTorch 里常见的写法是 nn.Sequential，键名是层的下标，ReLU 没有参数但占一个下标：
seq = nn.Module()
seq.net = nn.Sequential(nn.Linear(IN, HID), nn.ReLU(), nn.Linear(HID, OUT))
# 权重拷自主线模型
print("sequential keys:", list(seq.state_dict().keys()))
sequential keys: ['net.0.weight', 'net.0.bias', 'net.2.weight', 'net.2.bias']
burn 里对应的写法是 Vec<Linear>，激活函数单独放一个字段，路径是 net.0、net.1：
#[derive(Module, Debug)]
struct SeqMlp {
net: Vec<Linear>,
relu: Relu,
}
直接用 SafetensorsStore（挂了适配器）读：
Err
Tensor not found:
┌─ Tensor Loading Summary ─────────────────────────
│
│ ✓ Successfully applied: 2 tensors
│ ⊘ Skipped (filtered): 0 tensors
│ ✗ Missing in source: 2 tensors
│ ? Unused in target: 2 tensors
│ ! Errors: 0 errors
│
├─ Missing Tensors (requested by model but not found in source)
│
│ First 10 missing tensors:
│ • net.1.bias
│ Did you mean: 'net.2.bias'?
│ • net.1.weight
│ Did you mean: 'net.2.weight'?
│
├─ Unused Tensors (in source but not used by model)
│
│ First 10 unused tensors:
│ • net.2.bias
│ • net.2.weight
│
└───────────────────────────────────────────────────
Hint: Use `.allow_partial(true)` on the store to get an `ApplyResult` with structured missing/error info instead of a hard failure.
两种改法，都能加载、前向最大误差 1.49e-8：
- 把下标重编成连续的：store 上加
map_indices_contiguous(true)，net.2变成net.1。 - 不用
Vec，在加载时改名：with_key_remapping(r"^net\.0\.", "fc1.")和with_key_remapping(r"^net\.2\.", "fc2.")，读进原来的Mlp。
let mut st = SafetensorsStore::from_file(format!("{dir}/a1_seq.safetensors"))
.with_from_adapter(PyTorchToBurnAdapter)
.with_key_remapping(r"^net\.0\.", "fc1.")
.with_key_remapping(r"^net\.2\.", "fc2.");
改名规则怎么按顺序执行，见 burn 卷的存储与加载里 KeyRemapper 一节。
nn.Sequential 换成 torch.save(seq.state_dict()) 存的 .pt，用 PytorchStore::from_file 读进 SeqMlp，什么选项都不加，能不能成功？文件里的键名还是 net.0.*、net.2.*，SeqMlp 要的还是 net.0.*、net.1.*。先猜，再看本页末尾的答案。
dtype：float64 的文件
把主线模型 model.double() 之后再存（a1_f64.safetensors）。PyTorch 读回一个 f32 模型：
m = Mlp()
print(m.load_state_dict(load_file(f"{out}/a1_f64.safetensors")))
print("fc1.weight dtype after load:", m.fc1.weight.dtype)
<All keys matched successfully>
fc1.weight dtype after load: torch.float32
load_state_dict 把值拷进已有的参数，dtype 跟着参数走，f64 被转成 f32。burn 不一样：
let mut st = SafetensorsStore::from_file(format!("{dir}/a1_f64.safetensors"))
.with_from_adapter(PyTorchToBurnAdapter);
match model.load_from(&mut st) { /* ... */ }
println!("fc1.weight dtype {:?}", model.fc1.weight.val().dtype());
let logits = model.forward(a1_input(&device));
Ok
fc1.weight dtype F64
thread 'main' (8616552) panicked at .../burn-flex-0.22.1/src/ops/matmul.rs:84:5:
assertion `left == right` failed: matmul: dtype mismatch
left: F32
right: F64
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
加载成功，参数变成了 F64；输入是 F32，第一次 matmul 就 panic。burn-store 加载浮点参数时保留文件里的 dtype，只检查「浮点进浮点、整数进整数」（TargetDType 上面的注释写着「A loaded float or int parameter keeps whatever dtype the source declared」）；burn 也没有 dtype 提升，F32 和 F64 不能直接相乘（见 dtype 与设备）。要得到 F32 的模型，在适配器后面串一个 FloatCastAdapter：
let mut st = SafetensorsStore::from_file(format!("{dir}/a1_f64.safetensors"))
.with_from_adapter(PyTorchToBurnAdapter.chain(FloatCastAdapter::to(DType::F32)));
Ok
fc1.weight dtype F32
前向和 PyTorch 比，最大误差 1.49e-8。这里的 f64 值本来就是从 f32 扩上去的，转回 f32 没有损失。BF16 / F16 的文件也一样：加载返回 Ok、参数变成 BF16 / F16，前向在同一处报 dtype 不匹配，串同一个 FloatCastAdapter 就好，实测见常见坑加载 float64 权重后 matmul 报 dtype mismatch。
#4194「Fix dtype preservation when loading tensors in burn-store」（首次发布于 v0.20.0-pre.6）的提交说明：「Fixes a bug where tensor dtype was not preserved when loading from a snapshot, causing F16 tensors to be loaded as F32. Now uses Tensor::from_data_dtype to ensure the correct dtype is used.」之前加载时数据被转成后端的默认浮点类型，存成 F16 的模型读回来变成 F32。修好以后，文件里是什么 dtype，参数就是什么 dtype，上面 F64 文件的行为就是这么来的；applier.rs 里的测试 dtype_preservation_f64 专门钉住这一点。想按模型的 dtype 读，就像上面一样显式串一个 FloatCastAdapter。
动脑筋答案
Ok
┌─ Tensor Loading Summary ─────────────────────────
│
│ ✓ Successfully applied: 4 tensors
│ ⊘ Skipped (filtered): 0 tensors
│ ✗ Missing in source: 0 tensors
│ ? Unused in target: 0 tensors
│ ! Errors: 0 errors
│
└───────────────────────────────────────────────────
能成功，前向和 PyTorch 比最大误差 1.49e-8。同样的键名，SafetensorsStore 失败而 PytorchStore 成功，是因为两个 store 的 map_indices_contiguous 默认值不同：PytorchStore::map_indices_contiguous 的文档写着「(default: true)」，SafetensorsStore 的写着「(default: false)」。加这个选项的 #4150（v0.20.0-pre.6）在提交说明里写的就是「Enables this feature by default for PyTorchStore and adds configuration for both PyTorchStore and SafetensorsStore」。从 PyTorch 导出的 safetensors 要自己打开它。
小结
- PyTorch 存的 safetensors 读进 burn 要挂
PyTorchToBurnAdapter，Linear的权重从[out, in]转成[in, out]；读.pt的PytorchStore自带。主线 MLP 前向最大差 1.49e-8（容差 1e-6）。 - 不挂适配器：非方阵报
ShapeMismatch，方阵加载成功、结果是错的。导入后一律拿一个输入和 PyTorch 逐个数比。 PytorchStore能读 ZIP 格式、旧格式和torch.save(model)存的整个模型，不需要 Python、不执行 pickle 里的代码；PyTorch 2.15.0 的torch.load默认weights_only=True，整个模型反而读不了。nn.Sequential的下标有空洞：map_indices_contiguous(true)重编号（PytorchStore默认开，SafetensorsStore默认关），或者with_key_remapping改名。- burn 加载时保留文件的 dtype，F64 文件让参数变成 F64，前向时 dtype 不符 panic；串一个
FloatCastAdapter::to(DType::F32)。PyTorch 的load_state_dict跟着参数的 dtype 走。