大家好!我们很高兴向大家介绍mlx-rs
,这是一个让Rust开发者能够在Apple Silicon上实现高性能机器学习的项目。MLX框架是苹果专门为M系列芯片开发和优化的机器学习框架,而mlx-rs
则是基于苹果官方的C API构建的Rust绑定,让你能够在Mac上使用Rust进行高效的机器学习任务。
mlx-rs
旨在提供与官方Python和Swift版本相同的性能,并且API设计上也尽量保持一致,以便于大家快速上手。虽然mlx-rs
的文档还在不断完善中,但你可以参考苹果官方的Python和Swift版本文档来获取更多详细信息。
项目信息
- GitHub仓库: https://github.com/oxideai/mlx-rs
- 文档: 由于目前只能在M系列的macOS上编译,文档暂时托管在GitHub Pages上:https://oxideai.github.io/mlx-rs/mlx_rs/
- MSRV: 1.81
示例
让我们通过一个简单的MLP (Multi-Layer Perceptron) 示例来看看如何使用mlx-rs
:
use mlx_rs::{
error::Exception,
macros::ModuleParameters,
module::Module,
nn::{Linear, Relu, Sequential},
Array,
};
#[derive(Debug, ModuleParameters)]
pub struct Mlp {
#[param]
pub layers: Sequential,
}
impl Module<&Array> for Mlp {
type Error = Exception;
type Output = Array;
fn forward(&mut self, x: &Array) -> Result<Array, Self::Error> {
self.layers.forward(x)
}
fn training_mode(&mut self, mode: bool) {
self.layers.training_mode(mode);
}
}
impl Mlp {
pub fn new(
num_layers: usize,
input_dim: i32,
hidden_dim: i32,
output_dim: i32,
) -> Result<Self, Exception> {
let mut layers = Sequential::new();
// 添加输入层
layers = layers
.append(Linear::new(input_dim, hidden_dim)?)
.append(Relu);
// 添加隐藏层
for _ in 1..num_layers {
layers = layers
.append(Linear::new(hidden_dim, hidden_dim)?)
.append(Relu);
}
// 添加输出层
layers = layers.append(Linear::new(hidden_dim, output_dim)?);
Ok(Self { layers })
}
}
由于篇幅限制,更详细的训练代码可以参考GitHub上的示例: https://github.com/oxideai/mlx-rs/tree/main/examples/mnist
此外,我们还提供了一个从HuggingFace加载预训练权重的Mistral-7B-v0.1示例: https://github.com/oxideai/mlx-rs/tree/main/examples/mistral
1
共 0 条评论, 1 页
评论区
写评论还没有评论