Rust 深度学习库 Burn

一、概述

Burn 它是一个新的综合动态深度学习框架，使用 Rust 构建的，以极高的灵活性、计算效率和可移植性作为其主要目标。

Rust Burn 是一个以灵活性、高性能和易用性为核心设计原则工具，主打就是灵活性 、高性能 及易用性。

二、Rust Burn 拥有独特的特点

动态计算图： 它提供了灵活且可动态调整的计算图，这意味着开发者可以根据需要轻松地修改和优化数据流。

线程安全： Rust Burn 的数据结构设计考虑了线程安全性，确保在多线程环境下的稳定性和可靠性。

简化开发： 通过直观的抽象层，Rust Burn 旨在简化开发流程，让开发者能够更加专注于算法和逻辑的实现，而非底层细节。

快速性能： 无论是在模型训练还是推理阶段，Rust Burn 都能提供卓越的性能，显著提高工作效率。

多后端支持： 它支持多种后端实现，包括CPU和GPU，使得用户可以根据资源和需求选择合适的计算平台。

训练支持： Rust Burn 完全支持训练过程中的日志记录、度量跟踪和检查点保存，这些功能对于模型调试和优化至关重要。

开发者社区： 虽然Rust Burn 的开发者社区相对较小，但它活跃且充满热情，为用户和贡献者提供了一个友好的交流和协作环境。

三、快速入门

第一步： 确保已经安装 Rust 环境

第二步： 使用 Cargo 命令构建一个新工程

cargo new my_burn_app && cd my_burn_app/
1

第三步： 添加 Burn 作为依赖项，并添加启用 GPU 操作的 WGPU 后端功能

cargo add burn --features wgpu
1

第四步： 编译项目以安装 Burn

cargo build
1

成功安装 Burn 后，心动的时刻到了，可以开始使用这个强大的框架开发深度学习项目了。

四、编写示例

Rust 代码示例：

use burn::tensor::{Tensor, backend::Backend};

fn computation<B: Backend>() {
    // Create the device where to do the computation
    let device = Default::default();

    let tensor1: Tensor<B, 2> = Tensor::from_floats([[2., 3.], [4., 5.]], &device);
    let tensor2 = Tensor::ones_like(&tensor1);

    // Print the element-wise addition of the two tensors.
    println!("{:}", tensor1 + tensor2);
}

fn main() {
    computation::<burn::backend::Wgpu>();
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

编码完成，运行

cargo run
1

输出结果：

Tensor {
  data: [[3.0, 4.0], [5.0, 6.0]],
  shape: [2, 2],
  device: BestAvailable,
  backend: "wgpu",
  kind: "Float",
  dtype: "f32",
}
1
2
3
4
5
6
7
8

Rust Burn 选择了独特高性能，内存安全的语言来开发，也是符合现在未来趋势。尤其，现在 AI 大模型推动下，以前的很多东西都会被重塑，这是必然的。