从 PyTorch 到 burn

这一卷写给已经会写 PyTorch、想改用 burn（一个 Rust 深度学习框架）的读者。每章只做一件事，比如建一个张量、求一次梯度、训练一步、存一次模型，给出 PyTorch 和 burn 两份都能跑的程序，再贴两边的真实输出。先说两边概念上相同的部分，再说哪里不同、为什么不同。说「两边对得上」的地方，都给出实测的最大绝对误差和容差。

burn 内部怎么实现，这一卷不重讲。burn 源码解析那一卷讲过的机制，这里只给链接。这一卷讲的是从 PyTorch 过来，哪些写法要改、改成什么。

三条主线

全卷用同一个小模型串起来：两层 MLP（4→8→3，ReLU），交叉熵，SGD，学习率 0.1，权重按固定公式生成，两边逐个数对得上。每个内容页开头的「主线位置」说明这一页讲的是下面哪条线的哪一步。

• 张量主线：建张量 → 计算 → 把结果读回主机。对应 01 张量、02 所有权。
• 训练主线：forward → loss → backward → 优化器更新一步。对应 03 自动求导、04 模型、05 训练一步。
• 部署主线：保存 → 加载 → 不带自动求导地推理。对应 08 存取与迁移、09 部署。

每条主线的每一步落在哪一页，见怎么读这一卷。

flowchart TB
  intro[00 导读] --> tensor[01 张量]
  tensor --> own[02 所有权]
  own --> ad[03 自动求导]
  ad --> module[04 模型]
  module --> step[05 训练一步]
  step --> exec[06 执行方式]
  step --> data[07 数据加载]
  step --> store[08 存取与迁移]
  store --> deploy[09 部署]
  exec --> pit[10 常见坑]
  deploy --> pit
  pit --> hands[11 动手]

实验环境

每个对照都是一对程序：一个 .py 和一个 cargo crate，输出同样结构的 JSON 或逐行文本，再用一个几十行的 Python 脚本逐个数比较。平台是 M2 Max、macOS 27、rustc 1.99.0。burn 依赖锁定 =0.22.1，一律离线构建；每个实验写明开了哪些 feature、是 debug 还是 release。PyTorch 除了执行方式一部分对照 GPU（MPS 对 burn 的 Metal），都在 CPU 上跑。比速度的地方固定 PyTorch 的线程数，取多次运行的最小值，并写明当时机器的负载。

各部分

1. 00 导读：版本、对照程序怎么跑、「对得上」的标准；为什么要转、什么情况下不值得转。
2. 01 张量：创建与形状、dtype 与设备、归约与广播、整数运算、切片与索引。秩写在类型里，sum(dim) 保留维度，整数除法向零截断。
3. 02 所有权：运算按值消费张量，复用要 clone()；没有 add_ 这类带下划线的原地方法。
4. 03 自动求导：requires_grad 对 require_grad() 加 autodiff 设备；梯度不挂在张量上，而在 backward 返回的容器里；图只能用一次；不求导的几种写法。
5. 04 模型：nn.Module 对 #[derive(Module)]；遍历参数、冻结参数；训练和推理两种模式怎么切换。
6. 05 训练一步：一步训练逐项对照；优化器状态；训练循环。
7. 06 执行方式：逐个算子立即执行和攒起来再执行；计时要同步。
8. 07 数据加载：Dataset / DataLoader 对 burn 的 Dataset / Batcher / DataLoaderBuilder。
9. 08 存取与迁移：state_dict 对 burn-store；把 PyTorch 的权重读进 burn。
10. 09 部署：TorchScript / torch.export / ONNX 对单个 Rust 二进制。
11. 10 常见坑：从 PyTorch 习惯出发的二十六个坑，每条都有报错原文或静默的错误结果、最小复现和改法。
12. 11 动手：同一个小模型两边各训练到收敛，逐步对比 loss。

适合谁读

• 会写 PyTorch，能读懂 Rust 代码。这一卷不讲 Rust 语法，所有权只讲和张量有关的部分。
• 只想知道值不值得转：先读为什么要转，那一页同时写了不值得转的情况。
• 已经在写 burn，被 PyTorch 习惯绊住了：直接查常见坑，或者对应的那一章。