1000 行 rust 实现一个类似于 pytorch 的轻量级自动微分库

查看 5|回复 0
作者:StepfenShawn   
最近在学习神经网络原理,分享一下自己用 Rust 写的轻量级自动微分库 ferris-grad
Features
  • PyTorch 风格的 autograd 引擎
  • 核心代码不到 1000 行,纯 Rust 实现
  • 零外部依赖(仅使用 rust 标准库实现核心反向传播部分)

    包含什么?
    实现仅仅包含了 3 个文件构成:
    | 模块 | 功能 |
    |------|------|
    | scalar.rs | 标量计算图 + 反向传播 |
    | tensor.rs | 张量操作 |
    | nn.rs | 神经网络层 |
    能做什么?
  • 训练 MLP(~50 行代码)
  • 实现一个 mini GPT(~350 行代码)
  • ...

    代码示例:
    实现了一个 Pytorch 风格的 api:  
    use anyhow::Result;
    use ferris_grad::{Tensor, nn::Module};
    fn main() -> Result {
        let a = Tensor::from_vec(vec![1.0.into(), 2.0.into(), 3.0.into()], [3, 1].into())?;
        let b = Tensor::rand([3, 1].into())?;
        let c = &a * &b;
        println!("{}", c);
        Ok(())
    }
    本项目的灵感来源:
  • micrograd - Karpathy 的 python 微型 autograd 引擎  
  • microgpt - Karpathy 博客中的极简 GPT 实现

    GitHub 仓库地址: https://github.com/StepfenShawn/ferris-grad
    目前只能调用 CPU,后续研究一下如何在 GPU 上实现  
    欢迎 star 、提 issue 或者 PR !🦀
  • 您需要登录后才可以回帖 登录 | 立即注册

    返回顶部