性能
因为我们相信深度学习框架的目标是将计算转化为有用的智能,我们将性能作为 Burn 的核心支柱。 我们通过利用如下所述的多种优化技术来努力达到最高效率。
点击每个部分查看更多细节 👇
自动内核融合 💥
使用 Burn 意味着在任何后端都可以优化您的模型。 在可能的情况下,我们提供了一种自动和动态创建自定义内核的方法,该内核最小化不同内存空间之间的数据重新定位,当移动内存成为瓶颈时,这非常有用。
例如,您可以使用高级张量 API 编写自己的 GELU 激活函数(请参阅下面的 Rust 代码片段)。
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 中,反向传播实际上会改变每个张量参数的 grad 属性。 这不是线程安全的操作,因此需要较低级别的同步原语,详见 分布式训练。 请注意,这仍然非常快,但在不同的后端之间不兼容,并且相当难以实现。
智能内存管理 🦀
深度学习框架的主要作用之一是减少运行模型所需的内存量。 处理内存的简单方法是每个张量都有自己的内存空间,该空间在张量创建时分配,然后在张量超出作用域时释放。 但是,分配和释放数据非常耗时,因此通常需要内存池来实现良好的吞吐量。 Burn 提供了一种基础设施,可以轻松创建和选择针对后端的内存管理策略。 有关 Burn 中内存管理的更多详细信息,请参阅这篇博文。
Burn 的另一个非常重要的内存优化是我们会跟踪张量何时可以就地突变,只需使用所有权系统即可。 即使就其本身而言,这种内存优化相对较少,但在使用较大模型进行训练或运行推理时,它会累加并显著减少内存使用量。 有关更多信息,请参阅这篇关于张量处理的博客文章。
自动内核选择 🎯
一个优秀的深度学习框架应确保模型在所有硬件上都能顺畅运行。 但是,并非所有硬件在执行速度方面都有相同的行为。 例如,矩阵乘法内核可以使用许多不同的参数启动,这些参数对矩阵的大小和硬件非常敏感。 使用错误的配置可能会大幅降低执行速度(极端情况下甚至可能降低 10 倍甚至更多),因此选择正确的内核成为了优先事项。
使用我们自制的后端,我们会自动运行基准测试,并通过合理的缓存策略为当前硬件和矩阵大小选择最佳配置。
这会通过增加预热执行时间带来小的开销,但在几个前向和后向传递后会迅速稳定下来,从而在长远来看节省大量时间。 请注意,此功能不是强制性的,如果冷启动优先于优化的吞吐量,则可以禁用此功能。
硬件特定功能 🔥
众所周知,深度学习在很大程度上依赖于矩阵乘法作为其核心操作,因为这是如何建模全连接神经网络的。
越来越多的硬件制造商专门为了矩阵乘法工作负载优化他们的芯片。 例如,Nvidia 拥有其 Tensor Cores,而如今大多数手机都有 AI 专用芯片。 截至目前,我们支持使用 LibTorch 和 Candle 后端的 Tensor Cores,但尚不支持其他加速器。 我们希望 这个问题 在某些时候得到解决,以便为我们的 WGPU 后端提供支持。
自定义后端扩展 🎒
Burn 旨在成为最灵活的深度学习框架。 虽然保持与各种后端的兼容性至关重要,但 Burn 还提供了扩展后端实现的功能,以满足您的个人建模需求。
这种多功能性在很多方面是有利的,例如支持 flash attention 等自定义操作或手动为特定后端编写自己的内核以提高性能。 详见 Burn Book 🔥 中的本节。
训练与推理
使用 Burn,使整个深度学习工作流程变得容易,因为您可以使用符合人体工程学的仪表板来监控训练进度,并在从嵌入式设备到大型 GPU 集群的任何地方运行推理。
Burn 是从头开始构建的,考虑了训练和推理。 值得注意的是,与 PyTorch 等框架相比,Burn 简化了从训练到部署的过渡,消除了代码更改的需求。
点击以下部分展开 👇
训练仪表板 📈
如您在之前的视频中所见(点击图片!),基于 Ratatui crate 的新终端 UI 仪表板使用户能够轻松跟踪其训练,而无需连接到任何外部应用程序。
您可以实时可视化您的训练和验证指标更新,并使用箭头键分析任何注册指标的整体进展或最近的历史记录。在不崩溃的情况下休息训练循环,允许潜在检查点完全写入或重要代码片段完成而不间断 🛡
详情
ONNX 支持 🐫
ONNX(开放神经网络交换)是一种开放标准格式,可导出深度学习模型的架构和权重。
Burn 支持导入符合 ONNX 标准的模型,因此您可以轻松地将使用其他框架(如 TensorFlow 或 PyTorch)编写的模型移植到 Burn,以受益于我们框架提供的所有优势。
在 Burn Book 🔥 中进一步描述了我们的 ONNX 支持。
注意:这个 crate 正在积极开发中,目前支持了一套 有限的 ONNX 操作。
详情
导入 PyTorch 模型 🚚
支持将 PyTorch 模型权重加载到 Burn 的本地模型架构中,确保无缝集成。参见 Burn Book 🔥 中关于导入 PyTorch 的部分
详情
在浏览器中推断 🌐
我们的一些后端可以编译成 Web Assembly:Candle 和 NdArray 用于 CPU,WGPU 用于 GPU。这意味着您可以直接在浏览器中运行推断。我们提供了几个示例:
详情
嵌入式:no_std 支持 ⚙️
Burn 的核心组件支持 no_std。这意味着它可以在没有操作系统的裸机环境中运行,例如嵌入式设备。
截至目前,只有 NdArray 后端可以在 no_std 环境中使用。
详情
后端
Burn 力求在尽可能多的硬件上尽可能快,并有健壮的实现。我们相信这种灵活性对于现代需求至关重要,因为您可能会在云端训练模型,然后将其部署在客户的硬件上,这些硬件因用户而异。
与其他框架相比,Burn 支持多后端的方式非常不同。从设计上看,大部分代码都是基于通用后端特征的,这使得我们可以使用可替换的后端构建 Burn。这使得后端组成成为可能,通过像自动微分和自动内核融合等附加功能增强后端。
我们已经实现了许多后端,所有后端如下 👇
WGPU (WebGPU): 跨平台 GPU 后端 🌐
运行在任何 GPU 上的首选后端。
基于最受欢迎和支持最好的 Rust 图形库 WGPU,此后端自动针对 Vulkan、OpenGL、Metal、Direct X11/12 和 WebGPU,使用 WebGPU 着色语言 WGSL。它还可以编译成 Web Assembly,在浏览器中运行时利用 GPU,参见 此演示。有关此后端优势的更多信息,请参见 这个博客。
WGPU 后端是我们的第一个“内部后端”,这意味着我们对其实现细节有完全控制。它完全根据之前提到的性能特征进行了优化,因为它充当我们多种优化研究的实验场。
详情参见 WGPU 后端 README。
Candle:使用 Candle 绑定的后端 🕯
基于 Hugging Face 的 Candle,这是一款专注于性能和易用性的极简 ML 框架,此后端可在 CPU 上运行,并支持 Web Assembly 或使用 CUDA 的 Nvidia GPU。
详情参见 Candle 后端 README。
免责声明: 此后端尚未完全完成,但在某些推断环境中可以工作。
LibTorch:使用 LibTorch 绑定的后端 🎆
在深度学习领域,PyTorch 不需要介绍。此后端利用 PyTorch Rust 绑定,使您能够在 CPU、CUDA 和 Metal 上使用 LibTorch C++ 内核。
详情参见 LibTorch 后端 README。
NdArray:使用 NdArray 原语作为数据结构的后端 🦐
此 CPU 后端并非我们最快的后端,但提供了极端的便携性。
它是我们唯一支持 no_std 的后端。
详情参见 NdArray 后端 README。
Autodiff:为任何后端带来反向传播的后端装饰器 🔄
与上述后端不同,Autodiff 实际上是一个后端装饰器。这意味着它不能独立存在;它必须封装另一个后端。
简单的包裹一个基础后端并使用 Autodiff 会透明地为其增加自动微分支持,使您可以在模型上调用 backward。
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 后端 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 后端 README。
入门
刚听说 Burn?您来对地方了!继续阅读本节内容,希望您能很快上手。
Burn Book 🔥
要有效地与 Burn 一起工作,至关重要的是理解其关键组件和理念。这就是为什么我们强烈建议新用户阅读 The Burn Book 🔥 的前几部分。它提供了详细的示例和解释,涵盖了框架的每个方面,包括张量、模块和优化器等构建块,以及高级用法,如编写自己的 GPU 内核。
项目在不断发展,我们尽可能保持书籍与新添加的内容同步。然而,有时我们可能会遗漏一些细节,因此如果您看到奇怪的地方,请告诉我们!我们也欢迎 Pull Requests 😄
示例 🙏
<|vq_12199|> 让我们从一个代码片段开始吧,它展示了这个框架的直观使用方法!在下面的示例中,我们声明了一个带有一些参数的神经网络模块及其前馈过程。
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)
}
}
我们在仓库中有相当多的示例,展示了如何在不同场景中使用这个框架。
按照这本书进行:
- 基本工作流程 : 创建一个自定义的 CNN
Module
在 MNIST 数据集上进行训练和推理。 - 自定义训练循环 : 实现一个基本的训练循环,而不是使用
Learner
。 - 自定义 WGPU 内核 : 学习如何使用 WGPU 后端创建你自己的自定义操作。
其他示例:
- 自定义 CSV 数据集 : 实现一个数据集来解析 CSV 数据进行回归任务。
- 回归 : 在 CSV 数据集上训练一个简单的 MLP 用于回归任务。
- 自定义图像数据集 : 按照简单的文件夹结构在自定义图像数据集上训练一个简单的 CNN。
- 自定义渲染器 : 实现一个自定义渲染器来显示
Learner
的进度。 - 图像分类网页 : 使用 Burn、WGPU 和 WebAssembly 的图像分类网页演示。
- 网页上的 MNIST 推理 : 浏览器中的交互式 MNIST 推理演示。该演示可以在线查看。
- MNIST 训练 : 演示如何使用配置为记录指标和保持训练检查点的
Learner
训练一个自定义的Module
(MLP)。 - 命名张量 : 使用实验性的命名张量功能进行操作。
- ONNX 导入推理 : 导入一个预训练在 MNIST 上的 ONNX 模型,并在 Burn 上对样本图像进行推理。
- PyTorch 导入推理 : 导入一个预训练在 MNIST 上的 PyTorch 模型,并在 Burn 上对样本图像进行推理。
- 文本分类 : 在 AG News 或 DbPedia 数据集上训练一个文本分类 Transformer 模型。然后可以用训练好的模型对文本样本进行分类。
- 文本生成 : 在 DbPedia 数据集上训练一个文本生成 Transformer 模型。
为了获得更多实践见解,你可以克隆仓库并直接在你的计算机上运行任何这些示例!
预训练模型 🤖
我们保持一个更新和策划的模型和示例列表,这些都是用 Burn 构建的,详见 tracel-ai/models 仓库。
没有看到你想要的模型?不妨提交一个 issue,我们可能会优先考虑。使用 Burn 构建了一个模型并想分享吗?你也可以提交一个 Pull Request 并将你的模型添加到社区部分!
为什么使用 Rust 进行深度学习?🦀
深度学习是一种需要非常高层次抽象和极快执行时间的特殊软件。Rust 是这种用例的完美候选,因为它提供了零成本抽象,轻松创建神经网络模块,并通过内存的细粒度控制优化每一个细节。
一个框架在高层次上易于使用是很重要的,这样用户可以专注于 AI 领域的创新。然而,由于模型运行在很大程度上依赖于计算,性能不能被忽视。
到目前为止,解决这个问题的主流方案是在 Python 中提供 API,但依赖于绑定到底层语言如 C/C++。这降低了可移植性,增加了复杂性,并在研究者和工程师之间制造了摩擦。我们认为 Rust 对抽象的处理使其足够灵活,可以应对这两种语言之间的二分法。
Rust 还配有 Cargo 包管理器,使得从任何环境中构建、测试和部署变得非常简单,而在 Python 中通常这是一件痛苦的事。
尽管 Rust 一开始被认为是一种难学的语言,我们坚信它能在一些练习之后更快地构建出更加可靠、无虫的解决方案😅!
弃用说明
自0.14.0
起,张量数据的内部结构已更改。以前的Data
结构将被弃用,取而代之的是新的TensorData
结构,它通过将底层数据存储为字节并将数据类型保存为字段来提供更多的灵活性。如果你在代码中使用Data
,请务必切换到TensorData
。
加载旧版本的模型记录 ⚠️
如果你在尝试加载一个在旧版本中保存的模型记录,请确保启用 record-backward-compat
特性标志。
features = [..., "record-backward-compat"]
否则,记录将无法正确反序列化并且你会收到一条错误消息。这个错误也会指向后向兼容特性标志。
后向兼容性在加载记录时的反序列化过程中保持。因此,一旦你再次保存了记录,它将根据新的结构保存,并且你将不再需要后向兼容特性标志。
请注意,二进制格式不是向后兼容的。因此,你需要在旧版本中加载你的记录并将其保存为其他任何自描述记录格式(例如,使用 NamedMpkFileRecorder
)在使用带有 record-backward-compat
特性标志的新版本之前。
社区
如果你对这个项目感到兴奋,请不要犹豫,加入我们的Discord!我们尽力对每个人都表示欢迎,不论背景如何。你可以向社区提问并分享你所构建的东西!
贡献
在贡献之前,请花一点时间查看我们的行为准则。强烈推荐阅读架构概览,它解释了我们的一些架构决策。有关详细信息,请参阅我们的贡献指南。
状态
Burn 当前正在积极开发中,可能会有重大变化。虽然任何由此产生的问题可能都很容易解决,但在这一阶段没有任何保证。
许可证
Burn 根据 MIT 许可证和 Apache 许可证(版本 2.0)的条款分发。详见LICENSE-APACHE 和 LICENSE-MIT。提交 pull request 即表示同意这些许可条款。