跳到主要内容

Burn: 一个全面的动态深度学习框架

RustBurn深度学习框架
鱼雪

Burn 是一个新的全面的动态深度学习框架，以极高的灵活性、计算效率和可移植性为其主要目标。

性能​

我们相信深度学习框架的目标是将计算转化为有用的智能，因此我们将性能视为Burn的核心支柱。

我们努力通过利用下面描述的多种优化技术来实现最高效率。

自动内核融合 💥​

使用Burn意味着在任何后端对您的模型进行优化。

在可能的情况下，我们提供一种自动和动态创建自定义内核的方式， 最小化在不同内存空间之间的数据重定位，当内存迁移是瓶颈时非常有用。

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

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

然后，在运行时， 将自动为您的特定实现创建一个自定义低级内核，并将与手工编写的 GPU 实现相媲美。

内核由大约 60 行 WGSL WebGPU 着色语言组成，这是一种非常冗长的低级着色语言，您可能不希望用其来编写深度学习模型！

就目前而言，我们的融合策略仅为我们自己的 WGPU 后端实现，并且仅支持部分操作。

我们计划很快添加更多操作，并将这种技术扩展到未来的其他内部后端。

异步执行 ❤️‍🔥​

对由 Burn 团队从头开始开发的后端，使用了一种异步执行风格，这允许执行各种优化，例如先前提到的自动内核融合。 异步执行还确保框架的正常执行不会阻塞模型计算，这意味着框架的开销不会显著影响执行速度。 反之，模型中的强烈计算不会干扰框架的响应。 有关我们异步后端的更多信息，请参阅此博客文章。

线程安全的构建模块 🦞​

Burn通过利用Rust的所有权系统来强调线程安全。 使用Burn，每个模块都是其权重的所有者。

因此，可以将一个模块发送到另一个线程来计算梯度，然后将梯度发送到可以聚合它们的主线程，这样，您就可以进行多设备训练。

这与PyTorch的做法截然不同，PyTorch中的反向传播实际上会改变每个张量参数的梯度属性。 这是一个非线程安全的操作，因此需要较低级别的同步原语，请参考分布式训练。 请注意，这仍然非常快速，但不兼容不同的后端，并且实现起来相当困难。

智能内存管理 🦀​

深度学习框架的主要作用之一是减少运行模型所需的内存量。

处理内存的天真方式是每个张量都有自己的内存空间，在张量创建时分配，然后在张量超出范围时释放。

然而，分配和释放数据非常昂贵，因此通常需要内存池来实现良好的吞吐量。

Burn 提供了一个基础架构，可以轻松创建和选择后端的内存管理策略。

有关 Burn 中内存管理的更多详细信息，请参阅此博客文章。

Burn 的另一个非常重要的内存优化是，我们通过良好使用所有权系统，跟踪张量何时可以就地进行突变。

即使这在单独的内存优化方面相对较小，但在训练或运行较大模型进行推理时，会显着累积，并有助于进一步减少内存使用。

有关更多信息，请参阅有关张量处理的此博客文章。

自动内核选择🎯​

一个优秀的深度学习框架应该确保模型在所有硬件上运行顺畅。 然而，并非所有硬件在执行速度方面都表现相同。

例如，矩阵乘法内核可以通过许多不同的参数启动，这些参数对矩阵的大小和硬件非常敏感。 使用错误的配置可能会大幅降低执行速度（在极端情况下甚至会降低10倍或更多），因此选择正确的内核成为优先事项。

通过我们自制的后端，我们会自动运行基准测试，并选择适用于当前硬件和矩阵大小的最佳配置，配合合理的缓存策略。 这会稍微增加热身执行时间，但在几次前向和后向传递之后很快稳定下来，从长远来看节省大量时间。 请注意，这个功能不是强制性的，当冷启动优先于优化吞吐量时可以禁用。

硬件特定功能🔥​

深度学习主要依赖矩阵乘法作为其核心操作，因为这是全连接神经网络建模的方式。

越来越多的硬件制造商专门为矩阵乘法工作负载进行优化。 例如，Nvidia 拥有其 Tensor Cores，并且今天大多数手机都配备了人工智能专用芯片。

就目前而言，我们支持 Tensor Cores 与我们的 LibTorch 和 Candle 后端，但尚未支持其他加速器。 我们希望这个问题在某个时间点得到解决，以支持我们的 WGPU 后端。

自定义后端扩展 🎒​

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

这种多功能性在许多方面都具有优势，例如支持自定义操作，如快闪注意力，或手动编写特定后端的内核以增强性能。 请查看 Burn Book 🔥 中的此部分以获取更多详细信息。

训练和推理​

使用 Burn，整个深度学习工作流程变得轻松，因为您可以通过符合人体工程学的仪表板监控培训进展， 并在嵌入式设备到大型 GPU 集群中的任何位置运行推理。

Burn 从头开始构建时考虑了培训和推理。 值得一提的是，与 PyTorch 等框架相比，Burn 如何简化了从训练到部署的过度，消除了代码更改的需要。

培训仪表板 📈​

正如您在以前的视频中所看到 (单击图片！)， 一个基于 Ratatui 框架的新的终端 UI 仪表板使用户能够轻松地跟踪他们的培训，无需连接到任何外部应用程序。

您可以实时查看培训和验证指标的更新，并仅使用箭头键分析任何已注册指标的终身进展或最近历史。 在不崩溃的情况下中断培训循环，使潜在的检查点能够完全编写或重要的代码片段可以在没有干扰的情况下完成 🛡

ONNX 支持 🐫​

ONNX（Open Neural Network Exchange）是一种开放标准格式，可导出深度学习模型的架构和权重。

Burn 支持符合 ONNX 标准的模型的导入， 因此您可以轻松地将在另一个框架（如 TensorFlow 或 PyTorch）中编写的模型移植到 Burn， 以从我们的框架提供的所有优势中受益。

我们的 ONNX 支持在 Burn 书的这一部分中有进一步描述 🔥。

备注

此 crate 正在积极开发中，目前支持有限的 ONNX 运算符。

导入 PyTorch 模型 🚚​

支持将 PyTorch 模型权重加载到 Burn 的本地模型架构中，确保无缝集成。 参见 Burn 书 🔥 中有关导入 PyTorch 的部分。

浏览器内的推理 🌐​

我们的几个后端可以编译成 Web Assembly：Candle 和 NdArray 用于 CPU，WGPU 用于 GPU。

这意味着您可以直接在浏览器内运行推理。我们提供了几个示例：

  • MNIST， 您可以在其中绘制数字，然后一个小 convnet 试图找出它是什么！ 2️⃣ 7️⃣ 😰
  • 图像分类， 您可以上传图像并进行分类！ 🌄

嵌入式：不支持no_std ⚙️​

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

备注

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

后端​

Burn 旨在在尽可能多的硬件上尽可能快速，并具有强大的实现。 我们认为这种灵活性对于现代需求至关重要，在这里您可以在云中训练模型，然后在用户硬件上部署，这些硬件因用户而异。

与其他框架相比，Burn 对支持许多后端的方法有很大不同。 通过设计，大多数代码都是以 Backend 特质通用的，这使我们能够构建具有可互换后端的 Burn。 这使得构建后端变得可能，并通过增加额外功能，比如自动微分和自动核融合。

我们已经实现了很多后端，全部列在下面👇​

WGPU（WebGPU）：跨平台 GPU 后端 🌐​

可以运行在任何GPU上的首选后端。 基于最受欢迎和得到良好支持的Rust图形库 WGPU， 此后端自动以WebGPU着色语言WGSL为基础，针对Vulkan、OpenGL、Metal、Direct X11/12和WebGPU。 也可以编译成Web Assembly在浏览器中运行，同时利用GPU，可参见此演示。 有关此后端的更多信息，请参阅此博客。

WGPU后端是我们的第一个“内部后端”，这意味着我们完全控制其实现细节。 它充分优化了之前提到的性能特征，因为它是我们研究各种优化的实验场所。

有关更多详情，请查看WGPU后端README。

Candle：使用 Candle 绑定的后端 🕯​

基于 Hugging Face 的 Candle，这是一个专注于性能和易用性的极简主义 ML 框架， 该后端可以在支持 Web Assembly 的 CPU 上运行，也可以在支持 CUDA 的 Nvidia GPU 上运行。

有关更多详细信息， 请参阅 Candle 后端的 README。

备注

该后端尚未完全完成，但在某些情况下（例如推断）可以工作。

LibTorch：使用 LibTorch 绑定的后端 🎆​

在深度学习领域，PyTorch 无需介绍。 这个后端利用 PyTorch Rust 绑定，让您可以在 CPU、CUDA 和 Metal 上使用 LibTorch C++ 内核。

请查看 LibTorch 后端的 README 以了解更多详情。

NdArray：使用 NdArray 原始数据结构的后端 🦐​

这个 CPU 后端确实不是我们最快的后端，但提供了极强的可移植性。 这是我们唯一支持 no_std 的后端。

查看 NdArray Backend README 以获取更多详细信息。

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

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

简单地使用Autodiff将基础后端包装起来，可以透明地为其提供自动微分支持，从而可以在模型上调用反向传播。

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}"); 
}

请注意，对于不支持自动微分（用于推理）的后端，不可能在运行在此后端上的模型上调用backward这个方法， 因为这个方法只由Autodiff后端提供。

请参阅Autodiff Backend README 获取更多详细信息。

Fusion：为支持内核融合的后端带来内核融合的后端装饰器 💥​

此后端修饰符通过与内部后端支持的核融合增强后端。

请注意，您可以将此后端与其他后端修饰符（如 Autodiff）组合使用。 目前，仅 WGPU 后端支持融合内核。

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}");
}

值得注意的是，我们计划根据计算边界和内存边界操作实现自动梯度检查点， 这将与融合后端优雅地配合，使您的代码在训练期间运行得更快， 详细信息请参阅 Fusion Backend README。

入门​

刚听说过Burn吗？您来对地方了！继续阅读这一部分，希望您能很快上手。

The Burn Book 🔥​

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

备注

这个项目在不断发展，我们尽可能将这本书与新增内容保持最新。 然而，有时我们可能会漏掉一些细节，所以如果您发现有什么奇怪的地方， 请告诉我们！我们也很乐意接受Pull请求 😄

示例 🙏​

让我们从一个代码片段开始，展示这个框架使用起来有多直观！ 在以下内容中，我们声明一个神经网络模块，带有一些参数以及它的前向传递。

use burn::nn;
use burn::module::Module;
use burn::tensor::backend::Backend;

#[derive(Module, Debug)]
pub struct PositionWiseFeedForward<B: Backend> {
	linear_inner: nn::Linear<B>,
	linear_outer: nn::Linear<B>,
	dropout: nn::Dropout,
	gelu: nn::Gelu,
}

impl<B: Backend> PositionWiseFeedForward<B> {
	pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
		let x = self.linear_inner.forward(input);
		let x = self.gelu.forward(x);
		let x = self.dropout.forward(x);

		self.linear_outer.forward(x)
	}
}

我们在存储库中有相当多的示例， 展示了如何在不同情景中使用框架。 要获得更多实际见解，您可以克隆存储库并在计算机上直接运行任何一个示例！

预训练模型 🤖​

我们保持一个更新和精选的使用Burn构建的模型和示例列表，有关更多详细信息， 请参阅 tracel-ai/models 存储库。

找不到您想要的模型？不要犹豫，打开一个问题，我们可能会优先考虑。 使用Burn构建模型并想要分享吗？您还可以打开一个 Pull Request 并将您的模型添加到社区部分下！

为什么要使用 Rust 进行深度学习？🦀​

深度学习是一种特殊形式的软件，在这种软件中，你需要非常高级的抽象，以及极快的执行时间。

Rust 是这种用例的完美选择，因为它提供了零成本的抽象，可以轻松创建神经网络模块，并精细控制内存以优化每一个细节。

一个框架在高层次使用时易于使用，这样用户就可以专注于在人工智能领域进行创新。 但是，由于运行模型如此严重依赖计算，性能不能被忽视。

直到今天，解决这个问题的主流方法是在 Python 中提供 API，但依赖于诸如 C/C++ 等低级语言的绑定。 这降低了可移植性，增加了复杂性，并在研究人员和工程师之间产生摩擦。 我们认为 Rust 对于抽象的处理方式使其足够通用，可以解决这两种语言的困境。

Rust 还配备了 Cargo 软件包管理器，可以轻松构建、测试和部署任何环境的应用，通常在 Python 中很麻烦。

虽然 Rust 被认为是一种起初难以掌握的语言，但我们坚信经过一些练习后，它会带来更可靠、无 bug 的解决方案， 并且能够更快地构建（开玩笑😅）！

中年老码农,专注于自由职业与远程全栈开发
新文
关于站长和小站，聊聊经历与技术理念、职业发展、远程工作、接项目外包Rust 中的 `const` 与 `lazy_static` 的对比分析Python的包管理器: Rye 使用指南和注意事项Rust WebAssembly Hello World案例详解Rust Axum Web 应用通用化模版(01)