Burn Tensor

Burn Tensor库

该库提供多个Tensor实现，隐藏在易于使用的API后面，支持反向模式自动微分。

特性

• 灵活 ✨
• CPU + GPU 🙏
• 多线程 🚀
• 直观易用 😌
• 无全局状态 🚫
• 多个后端 🦾
• 反向模式自动微分 🔥

后端

目前实现了三个后端，并计划实现更多。

• 使用 tch-rs 的Pytorch后端
• 使用 ndarray 的100% Rust后端
• WGPU 后端
• Candle 后端
• 使用 tensorflow-rust 的Tensorflow后端
• 使用 RustCUDAtensorflow-rust 的CuDNN后端
• ...

自动微分

自动微分作为另一个Tensor后端实现，没有任何全局状态。这是可能的，因为我们跟踪每个操作执行顺序，并且只在计算梯度时创建tape。为此，每个操作创建一个新节点，该节点指向其父节点。因此，创建tape只需要一个简单高效的图遍历算法。

    let x = AutodiffTensor::from_tensor(x_ndarray);
    let y = ADtodiffTensor::from_tensor(y_ndarray);

    let z = x.matmul(&y);

    let grads = z.backward();

    let x_grad = x.grad(&grads);
    let y_grad = y.grad(&grads);

Cuda

要使用CUDA运行，请设置 TORCH_CUDA_VERSION=cu121。

备注

此crate可以独立使用，无需整个burn堆栈，并且仅用于较小的二进制文件选择后端。

功能标志

该包可以在不使用标准库的情况下使用（#![no_std]）并通过禁用默认的std特性来使用alloc。

• std - 启用标准库。
• burn-tensor-testgen - 启用生成张量测试的宏。