深度学习领域的Burn框架：兼顾灵活性与效率的完美之选

在当今快速发展的深度学习领域，选择一个既灵活又高效的框架至关重要。Burn 框架应运而生，它旨在提供一个不 的解决方案，兼顾灵活性、效率和可移植性，让开发者能够更好地构建和部署深度学习模型。

性能至上

Burn 的核心理念是将计算转化为有用的智能，因此性能是其最重要的支柱。为了实现最高的效率，Burn 采用了多种优化技术。

自动内核融合

使用 Burn 意味着你的模型可以在任何后端上进行优化。在可能的情况下，Burn 提供了一种自动且动态地创建自定义内核的方法，从而最大限度地减少不同内存空间之间的数据重新定位。这在移动内存成为瓶颈时非常有用。

例如，你可以使用高级张量 API 编写自己的 GELU 激活函数（见下面的 Rust 代码片段）。

fn gelu_custom<B: Backend, const D: usize>(x: Tensor) -> Tensor {    let x = x.clone() * ((x / SQRT_2).erf() + 1);    x / 2}

然后，在运行时，将自动为你的特定实现创建一个自定义的底层内核，其性能可以与手工制作的 GPU 实现相媲美。该内核包含大约 60 行 WGSL（ ）代码，这是一种极其冗长的底层着色器语言，你可能不想用它来编写深度学习模型！

异步执行

对于 Burn 的第一方后端，使用异步执行方式，这允许执行各种优化，例如前面提到的自动内核融合。

异步执行还确保框架的正常执行不会阻塞模型计算，这意味着框架开销不会显着影响执行速度。相反，模型中的密集计算不会干扰框架的响应能力。

线程安全构建块

Burn 通过利用 Rust 的所有权系统来强调线程安全。使用 Burn，每个模块都是其权重的拥有者。因此，可以将模块发送到另一个线程来计算梯度，然后将梯度发送到可以聚合它们的主线程，从而实现多设备训练。

这与 的做法截然不同，在 中，反向传播实际上会改变每个张量参数的grad属性。这不是线程安全的操作，因此需要较低级别的同步原语。请注意，这仍然非常快，但不兼容不同的后端，并且很难实现。

智能内存管理

深度学习框架的主要作用之一是减少运行模型所需的内存量。处理内存的简单方法是每个张量都有自己的内存空间，该空间在创建张量时分配，然后在张量超出范围时释放。但是，分配和释放数据的成本非常高，因此通常需要内存池才能实现良好的吞吐量。Burn 提供了一个基础架构，可以轻松地为后端创建和选择内存管理策略。

Burn 的另一个非常重要的内存优化是，我们通过使用所有权系统来跟踪何时可以在适当位置改变张量。即使它本身是一个相当小的内存优化，但在训练或使用更大的模型运行推理时，它会显着增加，并有助于进一步减少内存使用量。

自动内核选择

一个好的深度学习框架应确保模型在所有硬件上都能平稳运行。但是，并非所有硬件在执行速度方面都具有相同的行为。例如，可以使用许多不同的参数启动矩阵乘法内核，这些参数对矩阵的大小和硬件高度敏感。使用错误的配置可能会大大降低执行速度（在极端情况下甚至会降低 10 倍或更多），因此选择正确的内核成为当务之急。

使用我们自制的后端，我们会自动运行基准测试，并为当前硬件和矩阵大小选择最佳配置，并采用合理的缓存策略。

这会增加少量开销，从而增加预热执行时间，但在几次前向和后向传递后会快速稳定下来，从而从长远来看节省大量时间。请注意，此功能不是强制性的，并且可以在冷启动优先于优化吞吐量时禁用。

硬件特定功能

众所周知，深度学习主要依赖于矩阵乘法作为其核心运算，因为这是完全连接的神经网络的建模方式。

越来越多的硬件制造商专门针对矩阵乘法工作负载优化其芯片。例如， 有其 Cores，如今大多数手机都有 AI 专用芯片。截至目前，我们支持 Cores 与我们的 、、CUDA、Metal 和 WGPU/SPIR-V 后端，但不支持其他加速器。我们希望此问题能够尽快解决，以便为我们的 WGPU 后端提供支持。

自定义后端扩展

Burn 旨在成为最灵活的深度学习框架。虽然保持与各种后端的兼容性至关重要，但 Burn 还提供了扩展后端实现功能以满足你的个人建模需求的能力。

这种多功能性在许多方面都很有优势，例如支持自定义操作（如 flash ）或手动为特定后端编写自己的内核以提高性能。

多后端支持

Burn 致力于在尽可能多的硬件上以尽可能快的速度运行，并具有强大的实现。我们相信，这种灵活性对于现代需求至关重要，你可以在云中训练模型，然后在客户硬件上部署，这些硬件因用户而异。

「支持的后端」

后端

设备

类别

CUDA

GPU

第一方

ROCm

AMD GPU

第一方

Metal

Apple GPU

第一方

大多数 Linux 和 上的 GPU

第一方

Wgpu

大多数 GPU

第一方

大多数 CPU

第三方

大多数 GPU 和 CPU

第三方

, Apple GPU 和 CPU

第三方

与其他框架相比，Burn 在支持多个后端方面采用了一种非常不同的方法。通过设计，大多数代码都是通用的，可以通过 特征实现，这使我们能够使用可交换的后端构建 Burn。这使得可以组合后端，从而使用额外的功能（例如自动微分和自动内核融合）来增强它们。

：为任何后端带来反向传播的后端装饰器

与上述后端相反， 实际上是一个后端装饰器。这意味着它不能单独存在；它必须封装另一个后端。

用 包装基本后端的简单行为可以透明地为其配备自动微分支持，从而可以对你的模型调用 。

use burn::backend::{Autodiff, Wgpu};use burn::tensor::{Distribution, Tensor};
fn main() {    type Backend = Autodiff<Wgpu>;
    let x: Tensor<Backend, 2> = Tensor::random([32, 32], Distribution::Default);    let y: Tensor<Backend, 2> = Tensor::random([32, 32], Distribution::Default).require_grad();
    let tmp = x.clone() + y.clone();    let tmp = tmp.matmul(x);    let tmp = tmp.exp();
    let grads = tmp.backward();    let y_grad = y.grad(&grads).unwrap();    println!("{y_grad}");}

值得注意的是，不可能在不支持自动微分的后端（用于推理）上运行的模型上错误地调用 ，因为此方法仅由 后端提供。

：为所有第一方后端带来内核融合的后端装饰器

此后端装饰器使用内核融合来增强后端，前提是内部后端支持它。请注意，你可以将此后端与其他后端装饰器（例如 ）组合。目前，只有 WGPU 和 CUDA 后端支持融合内核。

use burn::backend::{Autodiff, Fusion, Wgpu};use burn::tensor::{Distribution, Tensor};
fn main() {    type Backend = Autodiff<Fusion<Wgpu>>;
    let x: Tensor<Backend, 2> = Tensor::random([32, 32], Distribution::Default);    let y: Tensor<Backend, 2> = Tensor::random([32, 32], Distribution::Default).require_grad();
    let tmp = x.clone() + y.clone();    let tmp = tmp.matmul(x);    let tmp = tmp.exp();
    let grads = tmp.backward();    let y_grad = y.grad(&grads).unwrap();    println!("{y_grad}");}

值得注意的是，我们计划实现基于计算绑定和内存绑定操作的自动梯度检查点，这将与融合后端完美配合，从而使你的代码在训练期间运行得更快。

(Beta)：将多个后端组合成一个后端的后端装饰器

如果例如你想在 CPU 上执行一些操作，而在 GPU 上执行其他操作，则该后端简化了硬件可操作性。

use burn::tensor::{Distribution, Tensor};use burn::backend::{    NdArray, Router, Wgpu, ndarray::NdArrayDevice, router::duo::MultiDevice, wgpu::WgpuDevice,};
fn main() {    type Backend = Router<(Wgpu, NdArray)>;
    let device_0 = MultiDevice::B1(WgpuDevice::DiscreteGpu(0));    let device_1 = MultiDevice::B2(NdArrayDevice::Cpu);
    let tensor_gpu =        Tensor::, 2>::random([3, 3], burn::tensor::Distribution::Default, &device_0);    let tensor_cpu =        Tensor::, 2>::random([3, 3], burn::tensor::Distribution::Default, &device_1);}

(Beta)：用于远程后端执行的后端装饰器，对于分布式计算很有用

该后端有两个部分：一个客户端和一个服务器。客户端通过网络将张量操作发送到远程计算后端。你可以在一行代码中使用任何第一方后端作为服务器：

fn main_server() {    // Start a server on port 3000.    burn::server::start::backend::Cuda>(Default::default(), 3000);}
fn main_client() {    // Create a client that communicate with the server on port 3000.    use burn::backend::{Autodiff, RemoteBackend};
    type Backend = Autodiff;
    let device = RemoteDevice::new("ws://localhost:3000");    let tensor_gpu =        Tensor::2>::random([3, 3], Distribution::Default, &device);}

训练与推理

借助 Burn，整个深度学习工作流程变得轻松，因为你可以使用符合人体工程学的仪表板监控训练进度，并在从嵌入式设备到大型 GPU 集群的任何地方运行推理。

Burn 从头开始构建时就考虑到了训练和推理。还值得注意的是，与 等框架相比，Burn 简化了从训练到部署的过渡，从而无需更改代码。

训练仪表板

正如你在之前的视频中看到的（点击图片！），基于 crate 的新终端 UI 仪表板允许用户轻松地跟踪他们的训练，而无需连接到任何外部应用程序。

你可以实时可视化训练和验证指标的更新，并仅使用箭头键分析任何已注册指标的终身进展或最近历史记录。从训练循环中断而不会崩溃，从而允许完全写入潜在的检查点或完整完成重要的代码

ONNX 支持

ONNX（开放神经网络交换）是一种开放标准格式，可导出深度学习模型的体系结构和权重。

Burn 支持导入遵循 ONNX 标准的模型，因此你可以轻松地将你在 或 等其他框架中编写的模型移植到 Burn，以从我们的框架提供的所有优势中受益。

导入 模型

支持将 模型权重加载到 Burn 的本机模型体系结构中，从而确保无缝集成。

浏览器中的推理

我们的几个后端可以编译为 Web ：用于 CPU 的 和 ，以及用于 GPU 的 WGPU。这意味着你可以直接在浏览器中运行推理。我们提供了几个这样的示例：

嵌入式：支持

Burn 的核心组件支持 。这意味着它可以运行在裸机环境中，例如没有操作系统的嵌入式设备。

截至目前，只有 后端可以在 环境中使用。

快速上手

想开始使用 Burn 了吗？你来对地方了！只需继续阅读本节，我们希望你可以很快上手。

Burn Book

要开始有效地使用 Burn，了解其关键组件和理念至关重要。这就是为什么我们强烈建议新用户阅读 Burn Book 的第一部分。它提供了详细的示例和解释，涵盖了框架的各个方面，包括张量、模块和优化器等构建块，一直到高级用法，例如编写自己的 GPU 内核。

示例

让我们从一个代码片段开始，该片段展示了框架的使用有多么直观！在下面，我们声明了一个带有某些参数的神经网络模块及其前向传递。

use burn::nn;use burn::module::Module;use burn::tensor::backend::Backend;
#[derive(Module, Debug)]pub struct PositionWiseFeedForward<B: Backend> {    linear_inner: nn::Linear,    linear_outer: nn::Linear,    dropout: nn::Dropout,    gelu: nn::Gelu,}
impl<B: Backend> PositionWiseFeedForward {    pub fn forwardD: usize>(&self, input: Tensor) -> Tensor {        let x = self.linear_inner.forward(input);        let x = self.gelu.forward(x);        let x = self.dropout.forward(x);
        self.linear_outer.forward(x)    }}