burn

Burn 是一个用 Rust 编写的张量库和深度学习框架，旨在用于数值计算、训练和推理。其明确目标是避免通常将 Python 训练栈与独立部署引擎分开的做法：同一模型代码可以先训练，然后直接在生产环境中运行，该项目认为这对设备端个性化和联邦学习很有用。 设计与人体工程学 - 类似 PyTorch 的 API，支持动态形状和图，同时对张量操作流进行 JIT 编译并执行自动内核融合。 - 强调增量编译：README 声称即使在 release 模式下，模型代码的修改也能在 5 秒内重新编译完成。 - 无需修改模型代码即可切换后端；多个后端可以在一个应用中共存，设备在运行时通过 `Device` 选择。 后端 - GPU：CUDA、ROCm、Metal、Vulkan、WebGPU 和 LibTorch（自 0.22.0 起已弃用）。各厂商覆盖情况不同：Nvidia（CUDA、Vulkan、WebGPU、LibTorch）、AMD（ROCm、Vulkan、WebGPU、LibTorch）、Apple（Metal、WebGPU、LibTorch）、Intel 和 Qualcomm（Vulkan、WebGPU）、Wasm（WebGPU）。 - CPU：CubeCL CPU 后端、Flex，以及适用于 x86 和 Arm 的 LibTorch；Flex 还覆盖 Wasm 和 no_std。 - 后端装饰器：Autodiff 为任何后端添加反向传播；Fusion 在受支持的情况下添加内核融合，并默认在第一方加速后端上启用；Remote（beta）支持通过 Iroh 或 WebSocket 进行客户端/服务器执行，用于分布式计算。 训练与推理 - 基于 Ratatui 构建的终端训练仪表盘实时显示训练和验证指标，支持方向键导航，并且可以在不崩溃的情况下中断训练循环。 - ONNX 模型可以通过 burn-onnx 导入，并使用 Burn API 转换为原生 Rust 代码，因此可以在任何 Burn 后端上运行；README 指出该 crate 正在积极开发中，算子集有限。 - 可以将 PyTorch 和 Safetensors 格式的权重加载到 Burn 定义的模型中。 - 推理可以通过 WebAssembly 在浏览器中运行，使用 Flex（CPU）或 WGPU（WebGPU），并提供 MNIST 和图像分类演示。 - 核心组件支持 no_std，适用于裸机嵌入式环境；目前只有 Flex 后端可在 no_std 中工作。 生态系统 - CubeCL：加速后端背后的 GPU 计算语言和编译器，可独立使用。 - burn-onnx：ONNX 导入；burn-store：权重保存/加载以及 PyTorch/Safetensors 导入。 - burn-vision、burn-rl、burn-dataset 分别用于视觉、强化学习和数据集。 - models 仓库包含预训练模型和示例；burn-bench 用于长期对后端进行基准测试。 - 社区 crate 列表涵盖数据加载（polars、arrow-rs、image、hf-hub）、分词与 NLP（tokenizers、rust-bert）、数值库（ndarray、nalgebra）、经典机器学习（linfa、smartcore）、推理运行时（candle、mistral.rs、ort、tract、wonx）、LLM/RAG 工具（rig、langchain-rust）、嵌入与向量搜索（fastembed、qdrant、lancedb）、计算机视觉（kornia-rs）、仿真（rapier）和可视化（rerun、plotters）。 入门 - Burn Book 是主要文档，涵盖张量、模块、优化器和自定义 GPU 内核。 - 示例包括基础 MNIST 工作流、自定义训练循环、自定义 WGPU 内核、CSV 和图像数据集、回归、自定义渲染器、浏览器推理演示、PyTorch 权重导入、文本分类和文本生成，以及 MNIST 上的 WGAN。 - 许可证为 MIT/Apache-2.0。