candle for Rust:Hugging Face 出品的轻量级机器学习框架
candle 是一个由 Hugging Face 团队用 Rust 开发的现代机器学习框架。与 PyTorch 或 TensorFlow 等重量级方案不同,candle 专注于高性能、内存安全和低资源消耗。它支持在 CPU 和 GPU(通过 CUDA 和 Vulkan)上运行神经网络模型,无需安装庞大的运行时环境。
为什么选择 candle?首先,它非常适合嵌入式系统、WebAssembly(WASM)和边缘设备,这些场景对二进制文件大小和加载速度有严格要求。其次,Rust 保证了线程安全且没有垃圾回收机制,为推理提供可预测的性能。第三,candle 支持从 PyTorch 格式(safetensors、bin)和 Hugging Face Hub 加载模型,是生产环境的绝佳选择。
截至本文撰写时,candle 正在积极开发中:支持 Transformer(BERT、GPT、LLaMA、Whisper)、卷积网络(ResNet、ViT)甚至生成模型(Stable Diffusion)。在本指南中,我们将介绍安装、核心方法,并编写一个完整的文本分类示例。
安装
candle 是一个 Rust 库,通过 Cargo 安装。在 Cargo.toml 中添加依赖:
[dependencies]
candle-core = "0.7"
candle-nn = "0.7"
candle-transformers = "0.7"
tokenizers = "0.21" # 用于分词器
anyhow = "1.0" # 用于错误处理
如果需要 CUDA 支持,使用 --features cuda 标志:
cargo add candle-core --features cuda
cargo add candle-nn --features cuda
cargo add candle-transformers --features cuda
如需 WebAssembly(WASM)支持,添加 --features wasm。
快速入门:最小工作示例
创建一个张量并执行简单操作——矩阵乘法:
use candle_core::{Device, Tensor};
fn main() -> Result<(), Box> {
// 在 CPU 上创建张量
let a = Tensor::new(&[[1.0f32, 2.0], [3.0, 4.0]], &Device::Cpu)?;
let b = Tensor::new(&[[5.0f32, 6.0], [7.0, 8.0]], &Device::Cpu)?;
// 矩阵乘法
let c = a.matmul(&b)?;
println!("结果: {:?}", c.to_vec2::()?);
Ok(())
}
这段代码将输出:
结果: [[19.0, 22.0], [43.0, 50.0]]
注意:所有操作都返回 Result,因此我们使用 ? 进行错误处理。
核心方法与函数
下面介绍 candle-core 和 candle-nn 中 10 个关键方法,它们将覆盖你 90% 的项目需求。
1. Tensor::new —— 创建张量
签名: pub fn new
说明: 从浮点数向量创建张量。形状通过 Shape 指定(例如 Shape::from((2, 3)))。
use candle_core::{Device, Tensor, Shape};
fn main() -> anyhow::Result<()> {
// 从向量创建 2x3 张量
let t = Tensor::new(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &Device::Cpu)?;
// 改变形状
let t_reshaped = t.reshape(&Shape::from((2, 3)))?;
println!("{:?}", t_reshaped.to_vec2::()?);
Ok(())
}
2. Tensor::matmul —— 矩阵乘法
签名: pub fn matmul(&self, other: &Tensor) -> Result<Tensor>
说明: 执行两个张量的矩阵乘法。支持批量广播。
use candle_core::{D