人工智能深度学习机器学习本地部署【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesnt compromise on flexibility, efficiency and portability.项目地址https://gitcode.com/GitHub_Trending/bu/burn点击查看免费下载导读本文围绕burn-dispatch这一 crate 展开它是 Burn 深度学习框架中“运行时后端选择”的枢纽——每个burn::tensor::Tensor背后的全局执行后端。读完本文你将掌握Dispatch 如何把同一份模型代码路由到 CPU / CUDA / wgpu / Flex / Remote 等多个后端并让它们并存Device构造函数与BURN_DEVICE环境变量如何决定运行时用哪个后端autodiff、梯度检查点与 kernel fusion 如何作为“装饰器”叠加以及#[backend_extension]如何把自定义算子接入 dispatch 层。全文以 crates/burn-dispatch/README.md 为骨架并结合仓库源码给出实现级证据。一、Dispatch 是什么Tensor 的“全局后端”Burn 的定位是“不牺牲灵活性、效率与可移植性的下一代张量库与深度学习框架”。要实现“同一份模型代码多种后端可跑”就需要一个能管理所有底层后端的统一执行层。burn-dispatch正是这个执行层。官方文档对该 crate 的定位是一句话点明Dispatchis the backend behind everyburn::tensor::Tensor.也就是说Burn 用户代码中写到的Tensor其底层的“后端类型”既不是Cpu也不是Cuda而是统一的Dispatch。它在编译期持有每一个被编译进应用的 backend 的张量 primitive张量句柄并在运行时把每个算子路由到拥有该张量的那个后端上执行。由此带来两个直接结果也是 Burn 区别于传统泛型后端框架的关键设计模型代码没有后端类型参数Tensor::2、Tensor::3这样的类型不再携带B: Backend泛型写模型时不需要TensorB, 2这种参数化签名一个程序可以同时使用多个后端只要把不同的Device传给不同的 TensorDispatch 就能让它们并行共存甚至在它们之间搬运数据。这一设计在 crates/burn/src/lib.rs 的 crate 文档中同样被强调——“Every enabled backend is available at runtime through aDeviceconstructor, and several can be used side by side”。二、一条 Tensor 操作的路由链路2.1 核心调用链Tensor → BridgeTensor → DispatchTensor → backend primitive原文档给出了 Dispatch 的数据流骨架Tensor - BridgeTensor - DispatchTensor - backend primitive逐层解读Tensor用户可见的高层 API位于 crates/burn-tensor携带形状、dtype 与DeviceBridgeTensor桥接层把高层 API 的调用转换为后端无关的操作描述DispatchTensor分发层张量定义在 crates/burn-dispatch/src/tensor.rs。它由两部分组成——kind: DispatchTensorKind真正指向哪个后端张量与autodiff: DispatchAutodiffContext携带该张量的自动微分 / 梯度检查点上下文backend primitive最终落到burn-cubecl、burn-flex、burn-remote等具体后端的张量原语上执行。关键数据结构是DispatchTensorKind——一个“每个已启用后端一个 variant”的枚举pub enum DispatchTensorKind { Cube(BackendTensorCube), // 所有 CubeCL 运行时共享这一个 variant Flex(BackendTensorFlex), NdArray(BackendTensorNdArray), LibTorch(BackendTensorLibTorch), Remote(BackendTensorRemote), Capture(BackendTensorCapture), Autodiff(BoxDispatchTensorKind), }其中BackendTensorB是对该后端 float / int / bool / quantized / autodiff 五种张量句柄的封装见 crates/burn-dispatch/src/tensor.rs。注意所有 CubeCL 运行时CUDA、ROCm、Metal、Vulkan、WebGPU、wgpu、CPU共用同一个Cubevariant运行时具体跑在哪个设备上由张量携带的Device决定——这正是后面“一个后端、多种运行时”设计的落点。2.2 算子如何被“转发”backend_dispatch过程宏在 crates/burn-dispatch/src/ops/tensor.rs 中FloatTensorOps、IntTensorOps、BoolTensorOps、QTensorOps、ModuleOps等 trait 的Dispatch实现几乎全部由#[backend_dispatch]属性宏生成例如#[backend_dispatch] impl FloatTensorOpsSelf for Dispatch { fn float_add(lhs: FloatTensorSelf, rhs: FloatTensorSelf) - FloatTensorSelf { B::float_add(lhs, rhs) } // ... 其余算子同理 }这里的B是宏生成的“路由结果”它根据输入张量的运行时后端 tag把B绑定到具体后端类型如Cube、Flex然后把B::float_add(..)转发到该后端。这一生成流水线定义在 crates/burn-backend-extension/src/lib.rs#[backend_dispatch] ─┐ ├─ ir::Operation ─ routing ─ generated enum dispatch #[backend_extension] ─┘ir描述张量输入、输出与后端调用与具体前端built-in dispatch 或用户 extension无关routing负责共享的后端选择、输入提取、调用与输出包装catalog唯一的“运行时后端清单”。需要特殊路由的方法如float_to_device的跨后端搬运用#[backend_dispatch(skip)]标记改由手写的float_to_device!宏生成完整的分发矩阵见 crates/burn-dispatch/src/ops/tensor.rs 与 crates/burn-dispatch/src/macros.rs。2.3 后端清单单一权威来源Dispatch 的路由宏backend_list!、distributed_backend_list!、backend_matrix!并不自行维护后端列表而是委托给burn-backend-extension的backend_catalog!过程宏见 crates/burn-dispatch/src/macros.rs。权威清单定义在 crates/burn-backend-extension/src/catalog.rs后端cfg 条件是否支持分布式distributedCubecube_backend是覆盖所有 CubeCL 运行时Flexfeature flex否NdArrayfeature ndarray否LibTorchfeature tch否Remotefeature remote是Capturefeature capture否单向转移有专用转移分支这份 catalog 同时派生出“分布式后端子集”与“转移矩阵”供所有生成路径与手写路径共用避免后端列表在多处维护导致漂移。三、支持的后端变体与 feature 组合3.1 后端总表原文档列出的后端变体如下Feature 列为在burn/burn-dispatch上启用的 Cargo feature变体Features底层 BackendCubecpu,cuda,metal,rocm,vulkan,webgpu,wgpuburn-cubecl覆盖所有 CubeCL 运行时Flexflexburn-flex纯 Rust CPURemoteremoteburn-remote另一台服务器上的设备Capturecaptureburn-capture只记录不执行NdArrayndarrayburn-ndarray已弃用LibTorchtchburn-tch已弃用两处要点README 与 crates/burn-dispatch/src/lib.rs 的模块文档一致确认所有 cubecl 后端共享同一个Cube后端feature 只决定“把哪些运行时编译进去”具体用哪个运行时由张量携带的 device 决定。可以同时启用多个运行时 feature它们共享唯一的DispatchDevice::Cubevariant。autodiff与fusion是装饰器autodiff可以把上面任意一个后端包装成 burn-autodifffusion则为 CubeCL 与 remote 后端开启 kernel fusion。feature 可以自由组合如flexautodifffusion。3.2 从源码看 feature 门控在 crates/burn-dispatch/src/backend.rs 中Dispatch的每个 Backend 方法都按DispatchDevice的 variant 展开 match 分支例如fn name(device: Self::Device) - String { let inner dispatch_device!(device, |device| B::name(device)); format!(dispatch{inner}) }dispatch_device!宏会为每个cfg满足条件的后端生成一个分支把B绑定为具体后端类型后执行同一段 body见 crates/burn-dispatch/src/macros.rs。若某个 variant 的 feature 未启用对应 match 分支直接通过#[cfg]消失——这就是“编译期裁剪后端、运行时选择后端”的实现机制。四、Device运行时的“后端选择器”4.1 Device 与 DispatchDevice 的关系用户面向的Device定义在 crates/burn-tensor/src/device.rs实际上是对DispatchDevice的类型擦除封装pub struct Device { blob: device_opaque::Opaque, // 类型擦除后的 DispatchDevice }而DispatchDevice定义在 crates/burn-dispatch/src/device.rs是一个“每个已启用后端一个 variant”的枚举pub enum DispatchDevice { Cube(CubeDevice), // 所有 CubeCL 运行时 Flex(FlexDevice), NdArray(NdArrayDevice), LibTorch(LibTorchDevice), Remote(RemoteDevice), Capture(CaptureDevice), Autodiff(AutodiffDevice), }Device提供了一组按 feature 门控的工厂方法见 crates/burn-tensor/src/device.rs每个方法把具体后端设备转换为DispatchDevice工厂方法Feature说明Device::cuda(0)cudaCUDA 设备整数或DeviceIndexDevice::rocm(0)rocmROCm/HIP 设备Device::cpu()cpuCubeCL CPU 运行时Device::wgpu(DeviceKind)wgpuwgpu自动选择图形 API 与编译器Device::metal(..)/vulkan(..)/webgpu(..)对应 feature分别钉死 Metal / Vulkan / WebGPUDevice::flex()flex纯 Rust CPU 后端Device::ndarray()/libtorch*(..)ndarray/tch已弃用Device::capture()capture非执行的图捕获设备底层转换关系在 crates/burn-dispatch/src/device.rs 中通过一组FromXxxDevice for DispatchDevice实现例如FromCudaDevice for DispatchDevice将其包装为DispatchDevice::Cube(CubeDevice::Cuda(..))——即 CUDA 设备最终也是落到Cubevariant 上印证了“一个 Cube 后端覆盖所有运行时”。4.2Device::default()的选择顺序与BURN_DEVICEDispatchDevice::default()的实现crates/burn-dispatch/src/device.rs遵循以下行为在std构建下若设置了环境变量BURN_DEVICE则按其值选择后端可接受值包括cuda、rocm、metal、vulkan、webgpu、wgpu、cpu、tch、remote、flex、ndarray未设置时按固定优先级选择第一个已启用的后端CUDA → Metal → ROCm → Vulkan → WebGPU → wgpu → CPU → LibTorch → Flex → Remote → NdArray若BURN_DEVICE指定了未知名称或未启用的 feature会直接panic!并给出配置提示若完全没有启用任何执行后端也会panic!提示启用flex、wgpu或cuda等后端 feature仅有capture时可通过Device::capture()记录图而不执行。注意一个细节默认顺序之所以“每个 feature 单独写死一个return”是为了避免CubeDevice::default()被 cargo feature 统一unification干扰——例如工作区同时构建了burn-cuda时即使本 crate 只启用了wgpuCubeDevice::default()也可能默认到 CUDA。逐个 feature 显式返回保证了“调用者未选择时的默认值”不随依赖图漂移源码注释对此有明确说明。4.3 设备枚举与运行时过滤Dispatch::enumerate(type_id)crates/burn-dispatch/src/backend.rs可以列出某类设备。对Cube而言枚举结果会经过cube_runtime_enabled过滤cargo 的 feature 统一可能让 cubecl 编译进比本 crate 请求更多的运行时比如 workspace 里有人启用了burn-cuda因此枚举必须按本 crate 自己的 feature过滤只交出当前构建真正启用的运行时设备见 crates/burn-dispatch/src/backend.rs。Dispatch::enumerate_cube(runtime)则进一步只列出指定 runtime 的设备。五、autodiff 与梯度检查点运行时上下文而非类型装饰Burn 0.22 的一个关键变化是autodiff 不再通过泛型AutodiffB出现在模型代码类型中而是作为DispatchDevice的运行时属性与DispatchTensor携带的上下文存在。5.1 Device 层.autodiff()与.gradient_checkpointing()Device::autodiff()把设备包装为DispatchDevice::Autodiff(AutodiffDevice)crates/burn-tensor/src/device.rsAutodiffDevice内含两个字段crates/burn-dispatch/src/device.rspub struct AutodiffDevice { inner: BoxDispatchDevice, checkpointing: GradientCheckpointingStrategy, }GradientCheckpointingStrategy只有两个取值crates/burn-dispatch/src/device.rs取值含义Disabled默认保留 autodiff 追踪但关闭梯度检查点Balanced反向传播时对选中激活值重算recompute以降低峰值内存代价是额外计算用户链式调用即可开启Device::default().autodiff().gradient_checkpointing()。如果对未开启 autodiff 的设备调用gradient_checkpointing()会 panic见 crates/burn-tensor/src/device.rs。Device::without_autodiff()则反向移除 autodiff 关联且操作幂等。5.2 Tensor 层DispatchAutodiffContext的合并规则DispatchTensor携带autodiff: DispatchAutodiffContextcrates/burn-dispatch/src/tensor.rs取值Disabled或Enabled(strategy)。一次算子调用中多个输入张量的上下文按以下规则合并源码中的merge方法Disabled与Enabled相遇disabled 张量被当作常数constant参与运算操作与输出使用 enabled 上下文两个Enabled相遇必须使用相同的梯度检查点策略否则 panic提示 “Gradient checkpointing strategy mismatch”。这一规则有专门的单元测试覆盖见 crates/burn-dispatch/src/backend.rs 底部的autodiff_context_tests模块例如disabled_and_enabled_float_contexts_merge_in_both_orders验证两种顺序合并结果一致fixed_arity_inputs_reject_mismatched_checkpointing_strategies验证策略不一致时 panicq_matmul_rejects_mismatched_checkpointing_strategies验证量化算子同样遵守该规则。5.3 运行时策略到编译期泛型的映射with_autodiff_backend!宏crates/burn-dispatch/src/macros.rs把运行时的GradientCheckpointingStrategy映射为编译期AutodiffB, C的具体泛型参数Balanced→BalancedCheckpointingDisabled→NoCheckpointing。这样 dispatch 层既能在运行时选择策略又能复用 burn-autodiff 的编译期优化。5.4 AutodiffBackend 的实现Dispatch实现了AutodiffBackendcrates/burn-dispatch/src/backend.rsbackward、grad、grad_remove、grad_replace、inner、from_inner等均按DispatchTensorKind的 variant 分发到对应后端的 autodiff 实现。其中Capture张量不支持 autodiff调用即 panic“Capture tensors do not support autodiff”inner/from_inner是“剥离 / 恢复 autodiff 上下文”的转换测试模块inner_transitions_clear_and_restore_context验证了上下文在转换中被正确清除与恢复grad返回的梯度张量上下文为Disabled“gradients are inner backend tensors”见gradients_are_inner_backend_tensors测试——即梯度本身是可脱离 autodiff 图的张量。六、跨后端张量转移to_deviceDispatch 的另一个核心职责是“管理跨后端张量转移”。在 crates/burn-dispatch/src/macros.rs 中to_device!与float_to_device!通过backend_matrix!生成所有后端两两组合的转移矩阵行为要点同后端转移走快路径直接把张量转移到该后端的新设备跨后端转移如Flex → Cube通过宿主内存中介float_transfer把源张量同步读回TensorData再用Dst::float_from_data在目标设备重建见 crates/burn-dispatch/src/ops/transfer.rs带 autodiff 的跨后端转移通过HostTransfer适配器同时记录前向与反向DifferentiableTransfer使梯度能沿转移路径反向传播见 crates/burn-dispatch/src/ops/transfer.rsCapture 是单向的普通后端张量可以被“物化”到 capture 设备成为捕获图中的 initializer但已捕获的张量没有真实数据不能搬回其他后端——转移矩阵中 capture 被刻意排除在交叉转移之外违规则 panic见 crates/burn-dispatch/src/macros.rs 及多处注释。转移相关的测试在 crates/burn-dispatch/src/ops/transfer.rs 底部例如recorded_transfer_uses_adapter_in_both_directions_without_replay验证AutodiffFlex → AutodiffNdArray转移时前向/反向适配器各被调用一次且梯度正确回到源设备。七、后端扩展#[backend_extension]自定义算子原文档指出后端扩展通过 burn-backend-extension 的#[backend_extension]宏为 dispatch 增加算子。仓库给出了完整示例示例项目examples/custom-cubecl-kernel演示如何编写自定义 CubeCL kernel并通过扩展 trait 接入 Dispatch。示例入口 examples/custom-cubecl-kernel/examples/custom-cubecl-kernel.rs 展示了使用方式——Device::default()取设备matmul_add_relu_custom调用自定义融合算子并与参考实现对比输出和梯度扩展 trait 的实现examples/custom-cubecl-kernel/src/lib.rs 中use burn::backend::{Dispatch, backend_extension, tensor::FloatTensor}说明自定义算子的签名直接写在Dispatch上。从实现原理看crates/burn-backend-extension/src/lib.rs#[backend_extension]为扩展 trait 生成Dispatch实现后端选择来自一个 routing tensor优先 float所有张量输入的 autodiff 上下文被合并disabled 输入视为常数enabled 输入必须共享同一梯度检查点策略struct / enum 输入通过ExtensionTypetrait定义在 crates/burn-core/src/backend.rs跨 dispatch 边界映射执行后端选择器selector包括Cube、Flex、NdArray、LibTorch、Remote注意Wgpu、Cuda等不是选择器——它们只是 CubeCL 运行时统一归入Cube自定义算子的 autodiff 支持需要为AutodiffB, C手写扩展 trait 的实现宏不自动生成反向传播。八、如何在自己的项目中使用 Dispatch实用清单虽然应用层不直接依赖burn-dispatch它是通过burn的 backend features 间接启用的但理解其工作机制有助于正确地配置你的项目在Cargo.toml中启用一个或多个后端 feature例如[dependencies] burn { version 0.22, features [wgpu] } # 或 [flex]、[cuda]、[cpu] 等Burn 默认不携带任何执行后端必须显式选择见 crates/burn/src/lib.rs 的 Feature Flags 说明。创建Device并在运行时选择后端use burn::prelude::*; let device Device::default(); // 按优先级自动选择 let cuda Device::cuda(0); // 显式指定需 cuda feature let flex Device::flex(); // 纯 Rust CPU需 flex feature let ad_device device.autodiff(); // 开启自动微分利用BURN_DEVICE环境变量在运行时切换默认后端便于同一份二进制在不同机器上跑不同硬件。让不同张量落在不同后端上并存Dispatch 允许Tensor::zeros([128, 128], cuda)与Tensor::zeros([128, 128], flex)共存于一个程序跨后端运算会自动触发to_device转移。需要自定义算子时参考 examples/custom-cubecl-kernel 与 crates/burn-backend-extension/README.md使用#[backend_extension]接入 Dispatch 路由。九、测试与验证Dispatch 正确性的保障仓库为 Dispatch 提供了多层测试验证上下文合并语义测试crates/burn-dispatch/src/backend.rs 的autodiff_context_tests需autodiffflexfeature覆盖策略传播、混用合并、不匹配 panic、int/bool/float 转换保持关联、量化算子上下文合并等设备 ID 往返测试crates/burn-dispatch/src/device.rs验证DispatchDevice::from_id(device.to_id())对 Remote 与 Capture 设备无损往返跨后端梯度传递测试crates/burn-dispatch/src/ops/transfer.rs验证AutodiffFlex → AutodiffNdArray前向/反向转移、分布式参数在不同后端上的拒绝与回环高层 API 集成测试crates/burn/tests 下的backend_extension_remote.rs、backend_extension_runtime.rs等验证扩展算子在不同运行时上的行为公开 API 测试crates/burn/tests/public_api_linalg.rs 等确保 dispatch 层的公开 API 形状稳定。十、小结burn-dispatch是 Burn 多后端架构的“中枢神经系统”它用DispatchDevice在运行时承载后端选择用DispatchTensor/DispatchTensorKind在运行时携带张量的后端归属与 autodiff 上下文用backend_dispatch过程宏把每个算子路由到正确后端用转移矩阵支撑跨后端数据搬运并用#[backend_extension]让第三方算子无缝接入。理解这一层就理解了为什么 Burn 的模型代码可以不带任何后端泛型参数却能在 CPU、GPU、Web 与远程设备之间自由切换与并存。相关进一步阅读后端扩展机制见 crates/burn-backend-extension/README.md后端实现见 crates/burn-cubecl/README.md 与 crates/burn-flex/README.md完整 API 见 crates/burn/src/lib.rs。赞分享人工智能深度学习机器学习本地部署【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesnt compromise on flexibility, efficiency and portability.项目地址https://gitcode.com/GitHub_Trending/bu/burn点击查看免费下载相关推荐深入Burn架构多后端支持与自动内核融合技术深入Burn架构多后端支持与自动内核融合技术 本文深入探讨了Burn深度学习框架的核心架构设计重点分析其多后端支持系统和自动内核融合技术。Burn通过精心设人工智能深度学习机器学习本地部署构建专业级物联网解决方案Arduino-ESP32开发框架的5大核心优势构建专业级物联网解决方案Arduino ESP32开发框架的5大核心优势 在嵌入式系统开发领域ESP32系列芯片以其卓越的Wi Fi/蓝牙双模通信能力和强大嵌入式物联网驱动开发DORA 运行时拆分dora-runtime-api SDK 与 OperatorRunner 多后端架构实战解析DORA 运行时拆分 dora runtime api SDK 与 OperatorRunner 多后端架构实战解析 导读 本文以 DORADataflow机器人人工智能ROS消息路由创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考