Basalt
纯Mojo🔥打造的机器学习框架
关于项目
Basalt是一个独立的机器学习框架,利用Mojo的强大功能。
正如Modular所讨论的,Mojo是面向AI开发未来的语言。它建立在MLIR技术之上,而不是现有的GCC和LLVM方法,Mojo的外观和感觉像Python代码,但性能更接近Rust或C++等语言。参数化函数和编译时参数允许静态编译图。静态图的存在使得更严格的性能优化成为可能。
尽管Basalt仍处于初期阶段,但它能够达到与成熟框架(如Pytorch)相当的速度。以下是当前基准测试的快照。但请持续关注,还有很大的改进空间,我们每天都在升级项目。
快速开始
亲自尝试基准测试:
mojo -I . examples/housing.mojo
mojo -I . examples/sin_estimate.mojo
mojo -I . examples/mnist.mojo
与替代的PyTorch实现进行比较:
确保在Python环境中安装python-requirements.txt
中的依赖项。
python examples/housing.py
python examples/sin_estimate.py
python examples/mnist.py
路线图
v0.1.0 ✅
- 改进矩阵乘法和卷积核
- 切换到自定义Tensor和TensorShape实现
- 改进基准测试和整体模型执行性能
- 添加性能分析和额外的性能测试
v0.2.0(进行中)
- 添加额外的运算符:Slice、(Un)Squeeze、Concat、Clip、Gather、Split、FMA等
- 更好的层支持和更多激活函数
- 图形子模块和图形连接
- 计算机视觉基准测试
长期目标
- 更好的并行化
- GPU支持
- 重新设计的数据加载器
- 自动调优和相关功能
- 图形编译优化
- 运算符融合
- ONNX / Max兼容性
贡献
Basalt由社区努力构建,依赖于您的专业知识和热情! 我们非常感谢小的修复和改进。如果您考虑更大的贡献,请随时联系我们以获得更流畅的Discord沟通渠道。如果您发现了bug或有功能想法,请使用我们的问题跟踪器。在创建新问题之前,请:
- 检查该问题是否已存在。如果已报告,您可以通过评论现有问题来贡献。
- 如果不存在,创建一个新问题并包含所有必要的细节以理解/重现问题或功能请求。
创建拉取请求
- Fork项目
- 创建您的功能分支
- 提交您的更改
- 推送到分支
- 打开拉取请求
一旦您推送了更改,请导航到GitHub上的fork。并针对原始basalt-org/basalt仓库创建拉取请求。
- 在创建PR之前,请确保它不会破坏任何单元测试。(例如
mojo run -I . test/test_ops.mojo
)- 引入新的大功能需要一个新的测试!
- 在拉取请求中,提供更改的详细描述以及为什么需要这些更改。链接任何相关问题。
- 如果有任何特定的测试或验证您的更改的说明,也请包括这些。
许可证
根据Apache 2.0许可证(带LLVM例外)分发。有关更多信息,请参阅LICENSE和LLVM许可证。