tracel-ai/burn

Burn – 一个以 Rust 为先的张量库和深度学习框架

它是什么 – Burn 是一个开源的 Rust crate，提供低级张量操作和高级深度学习 API。它的目标是让你编写相同的 Rust 代码来训练模型并在生产环境中运行该模型，从而省去了通常需要的 Python 到 ONNX 的导出步骤。

为什么它很重要 – 大多数研究框架（PyTorch、TensorFlow）都是以 Python 为中心的；训练好的模型随后会被导出到单独的运行时（ONNX、TensorRT 等）。Burn 统一了这些阶段：

• 类似于 PyTorch 的动态图人体工程学（形状可以在运行时更改）。
• JIT 编译和自动内核融合在保持动态图灵活性的同时提供了 GPU 级别的速度。
• Rust 中的增量编译使编辑-运行周期保持简短（即使在 release 模式下重新编译也 < 5 秒）。

关键组件

区域	Crate / 项目	它提供的内容
计算	CubeCL	一个基于 Rust 的 GPU 计算语言，可编译为 CUDA、ROCM、Metal、Vulkan 和 WebGPU。内核只需用 Rust 编写一次，即可在任何支持的后端上运行。
模型互操作	burn-onnx	导入 ONNX 模型，将其转换为原生的 Burn Rust 代码。
	burn-store	加载和保存权重（PyTorch .pt、Safetensors 等）。
领域库	burn-vision, burn-rl, burn-dataset	视觉操作、强化学习辅助工具和数据集实用程序。
预训练模型	models	使用 Burn 构建的可直接运行的示例集合。
工具	burn-bench	用于随时间比较后端性能的基准测试。

支持的后端 – Burn 通过插件架构可以在许多设备上运行：

• GPU – CUDA、ROCM、Metal、Vulkan、WebGPU（涵盖 Nvidia、AMD、Apple、Intel、Qualcomm 和浏览器）。
• CPU – 两个选项：
  • CubeCL CPU – 通过 LLVM 进行 JIT 编译的内核（支持 x86 和 ARM）。
  • Flex – 纯 Rust 渴望式后端，适用于 no_std、WebAssembly 和嵌入式目标。
• 后端可以通过 Device 抽象在运行时互换；你可以在同一个二进制文件中启用多个后端，并为每个张量选择设备。

特殊功能

• 自动微分 (Autodiff) – 反向传播是一个可以按设备开启的装饰器。张量仅在你请求时记录计算图，从而保持推理的轻量级。
• 内核融合 (Kernel fusion) – 自动将连续的操作合并为单个 GPU 内核，减少启动开销。
• 远程执行 (beta) – 客户端-服务器模型允许你使用加密的 Iroh 或简单的 WebSocket 传输将张量工作卸载到另一台机器（或云端）。
• 训练仪表板 – 一个终端 UI（使用 Ratatui 构建），显示实时指标，允许你安全地暂停训练，并且无需外部服务即可工作。
• ONNX 和 PyTorch 权重导入 – 无需重写即可将现有模型引入 Burn。
• WebAssembly 推理 – Flex (CPU) 和 WGPU 后端让你直接在浏览器中运行模型。

典型工作流程

1. 定义模型：使用 Burn 的 Module trait，就像在 PyTorch 中一样。
2. 创建 Device（例如 Device::cuda(0).autodiff() 或 Device::wgpu(Default::default()).autodiff()）。
3. 运行训练 – 张量自动记录梯度；调用 backward() 进行计算。
4. 保存权重：使用 burn-store 或导出为 ONNX 以实现互操作性。
5. 部署 – 切换 Device 到 CPU、WebGPU 或远程后端，无需更改模型代码。

谁应该使用它

• 研究人员：需要 Rust 的安全性和速度，但仍需要快速编辑-运行循环的人。
• 工程师：构建必须在异构硬件（云 GPU、边缘 CPU、浏览器）上运行的生产流水线，且希望使用单一代码库的人。
• 贡献者：有兴趣扩展 Rust AI 生态系统（新后端、领域 crate、模型库）的人。

入门 – 该 crate 已发布在 crates.io (burn = "<latest>")。文档位于 https://burn.dev/docs/burn，而“Burn Book”提供了关于张量、自动微分、训练循环和部署的分步教程。

---

以上所有信息均直接取自存储库的 README；未推断出任何额外功能。