< 返回版块

Mike Tang 发表于 2025-02-06 15:17

大家好!我们很高兴向大家介绍mlx-rs,这是一个让Rust开发者能够在Apple Silicon上实现高性能机器学习的项目。MLX框架是苹果专门为M系列芯片开发和优化的机器学习框架,而mlx-rs则是基于苹果官方的C API构建的Rust绑定,让你能够在Mac上使用Rust进行高效的机器学习任务。

mlx-rs旨在提供与官方Python和Swift版本相同的性能,并且API设计上也尽量保持一致,以便于大家快速上手。虽然mlx-rs的文档还在不断完善中,但你可以参考苹果官方的Python和Swift版本文档来获取更多详细信息。

项目信息

示例

让我们通过一个简单的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 页