dfdx: 在 Rust 中进行形状检查的深度学习

专注于易用性和安全性的 Rust 深度学习框架。

目前仍处于预 Alpha 阶段。接下来的几个版本预计会引入破坏性变更。

主要特性概览：

1. :fire: GPU 加速张量库，支持高达 6 维的形状！
2. 支持编译时和运行时均可变的维度形状。（例如 Tensor<(usize, Const<10>)> 和 Tensor<Rank2<5, 10>>）
3. 丰富的张量操作库（包括矩阵乘法、二维卷积等）。
   1. 所有张量操作在编译时都会进行形状和类型检查！！
4. 易用的神经网络构建模块（如全连接层、二维卷积层和 Transformer）。
5. 常用的深度学习优化器，如 SGD、Adam、AdamW、RMSprop 等。

dfdx 已发布在 crates.io 上！只需将以下内容添加到你的 Cargo.toml 文件中即可使用：

dfdx = "0.13.0"

更多文档请访问 docs.rs/dfdx。

[1] https://en.wikipedia.org/wiki/Automatic_differentiation#Reverse_accumulation

设计目标

1. 全面关注易用性（包括前端接口和内部实现）。
2. 尽可能在编译时完成所有检查（即如果存在错误，则无法编译）。
3. 最大化性能。
4. 尽量减少不安全代码[1]
5. 减少内部代码中 Rc<RefCell<T>> 的使用[2]

[1] 目前仅有的不安全调用是用于矩阵乘法的操作。

[2] 唯一使用 Arc 的地方是存储张量数据的部分。采用 Arc 而不是 Box 是为了减少张量克隆时的内存分配。

使用 CUDA 进行 GPU 加速

启用 cuda 功能即可开始使用 Cuda 设备！需要安装 NVIDIA 的 CUDA 工具包。更多信息请参阅 功能标志文档。

API 预览

更多详细信息请查看 examples/ 目录。

1. 👌 简单的神经网络 API，完全在编译时进行形状检查。

type Mlp = (
    (Linear<10, 32>, ReLU),
    (Linear<32, 32>, ReLU),
    (Linear<32, 2>, Tanh),
);

fn main() {
    let dev: Cuda = Default::default(); // 或 `Cpu`
    let mlp = dev.build_module::<Mlp, f32>();
    let x: Tensor<Rank1<10>, f32, Cpu> = dev.zeros();
    let y: Tensor<Rank1<2>, f32, Cpu> = mlp.forward(x);
    mlp.save("checkpoint.npz")?;
}

2. 📈 易用的优化器 API

type Model = ...
let mut model = dev.build_module::<Model, f32>();
let mut grads = model.alloc_grads();
let mut sgd = Sgd::new(&model, SgdConfig {
    lr: 1e-2,
    momentum: Some(Momentum::Nesterov(0.9))
});

let loss = ...
grads = loss.backward();

sgd.update(&mut model, &grads);

3. 💡 常量张量可以与普通的 Rust 数组相互转换

let t0: Tensor<Rank0, f32, _> = dev.tensor(0.0);
assert_eq!(t0.array(), &0.0);

let t1 /*: Tensor<Rank1<3>, f32, _>*/ = dev.tensor([1.0, 2.0, 3.0]);
assert_eq!(t1.array(), [1.0, 2.0, 3.0]);

let t2: Tensor<Rank2<2, 3>, f32, _> = dev.sample_normal();
assert_ne!(t2.array(), [[0.0; 3]; 2]);

有趣或值得注意的实现细节

模块

pub trait Module<Input> {
    type Output;
    fn forward(&self, input: Input) -> Self::Output;
}

通过这一灵活的特质，我们实现了：

1. 单个输入和批量输入（只需提供多个实现即可！）
2. 多个输入/输出（多头模块或 RNN）
3. 当梯度记录带存在与否时表现出不同的行为（这不同于其他库中的 .train() 和 .eval() 行为）！

元组表示前馈（即顺序）模块

由于我们可以为元组实现特质——据我所知，这是其他语言所不具备的——因此它们为按顺序执行模块提供了非常友好的前端接口。

// 不知道你为什么要这么做，但确实可以！
type Model = (ReLU, Sigmoid, Tanh);
let model = dev.build_module::<Model, f32>();

type Model = (Linear<10, 5>, Tanh)
let model = dev.build_module::<Model, f32>();

为 2 元组实现 Module 接口的方式如下：

impl<Input, A, B> Module<Input> for (A, B)
where
    Input: Tensor,
    A: Module<Input>,        // A 是一个接受 Input 的模块
    B: Module<A::Output>,    // B 是一个接受 A 输出的模块
{
    type Output = B::Output; // 输出是 B 的输出
    fn forward(&self, x: Input) -> Self::Output {
        let x = self.0.forward(x);
        let x = self.1.forward(x);
        x
    }
}

我们已经为最多 6 元素的元组实现了模块，而且你可以任意嵌套它们！

不使用 Rc<RefCells<T>> —— 梯度记录带并不保存在引用单元中！

其他实现可能会直接将梯度记录带的引用存储在张量上，这就需要频繁地修改张量或大量使用 Rc/RefCell。

而我们则找到了一种优雅的方法来避免这种情况，将引用和动态借用检查降至零！

由于所有操作都恰好产生一个子节点，我们可以始终将梯度记录带移动到最后一个操作的子节点上。此外，任何模型参数（即所有张量）都不会拥有梯度记录带，因为它们永远不会成为任何操作的结果。这意味着我们可以准确地知道哪个张量拥有梯度记录带，而拥有它的张量总是中间结果，无需在整个梯度计算过程中一直保留。

这一切共同赋予用户对哪些张量会被记录在梯度记录带上前所未有的控制力和精确度！

一个高级用例要求在计算图中多次重复使用同一个张量。这可以通过克隆该张量并手动移动梯度记录带来实现。

类型检查的反向传播

简而言之：如果你忘记调用 trace() 或 traced()，程序将无法编译！

-let pred = module.forward(x);
+let pred = module.forward(x.traced(grads));
let loss = (y - pred).square().mean();
let gradients = loss.backward();

由于我们确切地知道哪个张量拥有梯度记录带，因此我们可以要求传递给 .backward() 方法的张量必须拥有梯度记录带！更进一步，我们还可以要求将其移入 .backward() 方法中，以便销毁梯度记录带并计算梯度！

所有这些都可以在编译时进行检查 🎉

📄 与 PyTorch 对照验证

所有函数和操作都经过测试，确保其行为与 PyTorch 中类似代码的行为一致。

许可证

双重许可，以兼容 Rust 项目。

根据您的选择，本软件可根据 Apache License, Version 2.0 http://www.apache.org/licenses/LICENSE-2.0 或 MIT 许可证 http://opensource.org/licenses/MIT 进行授权。除非按照这些条款的规定，否则不得复制、修改或分发本文件。

dfdx 快速上手指南

dfdx 是一个专注于人体工学与安全性的 Rust 深度学习库。其核心特色是在编译期进行严格的形状检查（Shape Checked），确保张量维度不匹配时无法通过编译，同时支持 GPU 加速。

注意：目前该项目处于 pre-alpha 阶段，后续版本可能会包含破坏性更新。

环境准备

系统要求

• 操作系统：Linux, macOS, Windows
• 编译器：最新稳定版 Rust (rustc)
• 构建工具：cargo

前置依赖

• CPU 模式：无额外依赖，开箱即用。
• GPU 模式 (CUDA)：
  • 需安装 NVIDIA CUDA Toolkit。
  • 确保显卡驱动支持对应的 CUDA 版本。

安装步骤

在你的 Rust 项目目录中，编辑 Cargo.toml 文件添加依赖。

1. 基础安装 (CPU)

[dependencies]
dfdx = "0.13.0"

2. 启用 GPU 加速 (可选)

若需使用 CUDA 加速，请启用 cuda 特性：

[dependencies]
dfdx = { version = "0.13.0", features = ["cuda"] }

国内加速建议：如果下载 crates.io 依赖较慢，可配置国内镜像源。在 ~/.cargo/config.toml 中添加：

[source.crates-io]
replace-with = 'ustc'

[source.ustc]
registry = "sparse+https://mirrors.ustc.edu.cn/crates.io-index/"

基本使用

dfdx 的核心优势在于利用 Rust 类型系统在编译期检查神经网络结构。以下示例展示了一个简单的多层感知机 (MLP) 的构建、前向传播及优化器使用流程。

1. 定义模型与前向传播

通过元组 (Tuple) 定义序列模型，利用泛型指定输入输出维度，编译器将自动检查形状匹配。

use dfdx::prelude::*;
use dfdx::nn::{Linear, ReLU, Tanh};

// 定义模型结构：(Input:10 -> Hidden:32) -> ReLU -> (Hidden:32 -> Hidden:32) -> ReLU -> (Hidden:32 -> Output:2) -> Tanh
type Mlp = (
    (Linear<10, 32>, ReLU),
    (Linear<32, 32>, ReLU),
    (Linear<32, 2>, Tanh),
);

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 选择设备：Cuda (需开启 cuda 特性) 或 Cpu
    let dev: Cpu = Default::default(); 
    
    // 构建模型
    let mlp = dev.build_module::<Mlp, f32>();
    
    // 创建输入张量 (形状为 [10])
    let x: Tensor<Rank1<10>, f32, Cpu> = dev.zeros();
    
    // 前向传播 (形状在编译期已检查，输出形状自动推导为 [2])
    let y: Tensor<Rank1<2>, f32, Cpu> = mlp.forward(x);
    
    println!("Output shape is checked at compile time!");
    
    // 保存模型
    // mlp.save("checkpoint.npz")?; 
    
    Ok(())
}

2. 使用优化器训练

dfdx 提供了类型安全的反向传播机制。必须显式调用 .traced() 将张量加入梯度磁带，否则无法编译。

use dfdx::prelude::*;
use dfdx::optim::{Sgd, SgdConfig, Momentum};

// 假设 Model 已定义
type Model = (Linear<10, 2>, ReLU); 

fn train_step() -> Result<(), Box<dyn std::error::Error>> {
    let dev: Cpu = Default::default();
    let mut model = dev.build_module::<Model, f32>();
    
    // 分配梯度存储
    let mut grads = model.alloc_grads();
    
    // 初始化优化器 (例如带动量的 SGD)
    let mut sgd = Sgd::new(&model, SgdConfig {
        lr: 1e-2,
        momentum: Some(Momentum::Nesterov(0.9))
    });

    // 模拟数据
    let x: Tensor<Rank1<10>, f32, _> = dev.sample_normal();
    let y: Tensor<Rank1<2>, f32, _> = dev.sample_normal();

    // 关键步骤：必须使用 .traced(grads) 包裹输入，才能进行反向传播
    // 如果忘记此步，代码将无法通过编译
    let pred = model.forward(x.traced(grads));
    
    // 计算损失
    let loss = (y - pred).square().mean();
    
    // 反向传播：消耗 loss 张量并生成梯度
    let gradients = loss.backward();
    
    // 更新模型参数
    sgd.update(&mut model, &gradients);

    Ok(())
}

3. 张量与数组互转

对于编译期已知形状的张量（Const tensors），可以直接与 Rust 原生数组互换。

let dev: Cpu = Default::default();

// 标量
let t0: Tensor<Rank0, f32, _> = dev.tensor(0.0);
assert_eq!(t0.array(), &0.0);

// 一维向量
let t1: Tensor<Rank1<3>, f32, _> = dev.tensor([1.0, 2.0, 3.0]);
assert_eq!(t1.array(), [1.0, 2.0, 3.0]);

// 二维矩阵
let t2: Tensor<Rank2<2, 3>, f32, _> = dev.sample_normal();
// 获取原生数组 [[f32; 3]; 2]
let arr = t2.array();