跳到主要内容

Burn:一个灵活高效的 Rust 深度学习框架,支持自动内核融合与多后端集成

鱼雪

Burn 是一个新型的、全面的 动态深度学习框架,使用 Rust 构建,目标是提供极高的 灵活性计算效率可移植性。通过自动内核融合、异步执行和多硬件后端的支持,Burn 为开发者提供了前所未有的高性能深度学习体验。

Burn 的主要特性

  • 自动内核融合:通过动态生成 GPU 内核,提升模型执行效率。
  • 异步执行架构:保证框架与计算的独立性,提升系统性能。
  • 多后端集成:支持多种硬件架构,跨平台高效执行。
  • 线程安全与智能内存管理:利用 Rust 的所有权系统,确保训练过程中的安全与效率。

性能优化:自动内核融合与异步执行

自动内核融合 💥

Burn 提供自动内核融合,意味着在任何后端都能对模型进行优化。在可能的情况下,Burn 自动创建自定义低级内核,特别适用于需要最小化内存迁移的场景,从而显著提升模型的计算效率。

以下是一个 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
}

在运行时,Burn 会为这个自定义实现自动创建内核,与手工编写的 GPU 实现性能相媲美。目前,自动内核融合仅支持 WGPU 后端,未来会扩展到其他后端。

异步执行 ❤️‍🔥

Burn 的后端采用 异步执行风格,这使得模型的计算不会阻塞框架的正常执行,确保系统的高响应性。这种异步架构还为自动内核融合等优化技术提供了支持,进一步提高执行效率。更多关于异步执行的信息,请访问 Burn 博客

线程安全的模块 🦞

Burn 利用 Rust 的所有权系统,使每个模块成为其权重的唯一所有者。这样可以安全地将模块发送到不同线程进行计算,适合多设备训练,避免了传统框架中的同步问题。

智能内存管理 🦀

内存管理 是深度学习框架的关键之一。Burn 使用内存池减少频繁的内存分配和释放,从而提升吞吐量,并通过跟踪张量的所有权在适当时候进行就地突变,进一步减少内存使用。有关内存管理的更多详细信息,请参考 Burn 博客

自动内核选择 🎯

Burn 会自动运行基准测试并为硬件选择最佳内核配置,以确保在所有硬件上高效执行。这稍微增加了热身时间,但会在几次前向和后向传递后稳定下来,大大提高长期执行效率。

硬件特定优化 🔥

Burn 支持多种硬件特性,例如 Nvidia Tensor Cores。目前,Tensor Cores 支持通过 LibTorch 和 Candle 后端进行优化,但还未支持其他加速器。我们计划在未来为 WGPU 后端引入类似的硬件支持。

扩展与灵活性:后端扩展与自定义实现

Burn 提供了丰富的扩展能力,使开发者能够轻松地为模型添加自定义操作或编写特定后端的内核。例如,可以手动编写快闪注意力的实现来提升性能。更多关于后端扩展的信息,请参考 Burn Book

Burn 的深度学习工作流程

培训与推理

Burn 从一开始就考虑了深度学习的培训和推理过程,提供了符合人体工程学的仪表板来监控训练进展,并能够在从嵌入式设备到大型 GPU 集群的任意设备上运行推理。

  • 培训仪表板 📈:Burn 提供基于 Ratatui 框架的终端 UI 仪表板,用户可以方便地实时跟踪训练和验证指标。
  • ONNX 支持 🐫:Burn 支持导入符合 ONNX 标准的模型,便于将其他框架(如 TensorFlow 或 PyTorch)的模型移植到 Burn。
  • 浏览器推理 🌐:通过编译为 Web Assembly,Burn 支持直接在浏览器内运行推理。您可以查看以下示例:

多后端支持

Burn 致力于成为一个 跨平台、支持多后端的深度学习框架,旨在满足不同硬件和开发需求的灵活性。

支持的后端

  1. WGPU 后端 🌐:基于 Rust 图形库 WGPU,支持跨平台 GPU 加速,适用于从 Vulkan 到 WebGPU 等多种环境。更多信息,请参考 WGPU 后端 README
  2. Candle 后端 🕯:基于 Hugging Face 的 Candle,支持 Web Assembly 和 CUDA,适合极简的高性能模型。更多信息,请参考 Candle 后端 README
  3. LibTorch 后端 🎆:基于 PyTorch 的 Rust 绑定,支持 CPU、CUDA 和 Metal,加速深度学习模型训练和推理。更多信息,请参考 LibTorch 后端 README
  4. NdArray 后端 🦐:轻量级 CPU 后端,唯一支持 no_std 的后端,可在无操作系统环境下运行。更多信息,请参考 NdArray 后端 README
  5. Autodiff 后端 🔄:作为后端修饰器,为任何基础后端增加自动微分支持,用于模型训练。更多信息,请参考 Autodiff 后端 README
  6. Fusion 后端 💥:为支持内核融合的后端增加内核融合特性,目前仅 WGPU 后端支持融合。更多信息,请参考 Fusion 后端 README

入门指南

The Burn Book 🔥

要有效地开始使用 Burn,建议阅读 The Burn Book 的前几章,了解 Burn 的关键组成部分和哲学,涵盖了张量、模块和优化器等构建模块的详细示例。

示例代码 🙏

以下是一个简单的代码片段,展示 Burn 的模块声明和前向传播实现:

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。

为什么选择 Rust 进行深度学习 🦀

Rust 提供了 零成本抽象 和强大的 内存管理,是深度学习开发的理想选择。Rust 提供高性能和安全性,同时使用 Cargo 作为包管理器,可以轻松地构建、测试和部署应用程序。

Rust 的学习曲线可能较为陡峭,但一旦掌握,它可以提供更可靠、无 bug 的解决方案,极大地提升开发效率和代码质量。

结论

Burn 是一个功能强大的 Rust 深度学习框架,支持自动内核融合、异步执行、多后端集成,适用于从嵌入式设备到大型 GPU 集群的各种场景。如果您对深度学习的性能和灵活性有高要求,并且希望探索 Rust 的强大能力,那么 Burn 会是一个绝佳的选择。