Apple机器学习框架MLX在Node.js环境终于可用了!

前有科技后进阶 2024-06-13 08:49:42

大家好,很高兴又见面了,我是"高级前端‬进阶‬",由我带着大家一起关注前端前沿、深入前端底层技术,大家一起进步,也欢迎大家关注、点赞、收藏、转发,您的支持是我不断创作的动力。

什么是 node-mlx

MLX(Github 有 15k 的 star)是 Apple 芯片上用于机器学习研究的框架,由 Apple 机器学习研究团队研发。

MLX 的一些主要功能包括:

熟悉的 API:MLX 遵循 NumPy 的 Python API,具有功能齐全的 C++、C 和 Swift API,与 Python API 非常相似。同时, MLX 拥有更高级别的软件包,例如: mlx.nn 和 mlx.optimizers,API 遵循 PyTorch 以简化构建更复杂的模型。可组合函数转换:MLX 支持用于自动微分、自动向量化和计算图优化的可组合函数转换惰性计算:MLX 中的计算是惰性的,数组仅在需要时才会具体化。动态图构建:MLX 中的计算图是动态构建的,更改函数参数的形状 (Shape) 不会触发缓慢的编译,并且调试简单直观。支持多设备:操作可以在任何支持的设备上运行(当前是 CPU 和 GPU)统一内存:与其他框架的显著区别是统一内存模型,MLX 中的数组位于共享内存中。 MLX 阵列上的操作可以在任何支持的设备类型上执行,而无需传输数据。

而本文的主角 node-mlx 是基于 MLX 的 Node.js 机器学习框架,且与 Apple 官方无关。目前 GPU 支持采用 Apple Silicon 的 Mac。CPU 支持:

x64 Macx64/arm64 Linux目前尚不支持 Windows,但后续会继续适配 MLX

需要注意的是, MLX 没有计划支持 Apple Silicon 以外的 GPU,要使用 NVIDIA GPU 进行计算则必须使用 TensorFlow.js,或者等待有人将 PyTorch 移植到 Node.js。

如何使用 node-mlx

node-mlx 需要安装依赖,使用起来也非常简单,比如下面的示例:

import mlx from '@frost-beta/mlx';const {core: mx, nn} = mlx;const model = new nn.Sequential( new nn.Sequential(new nn.Linear(2, 10), nn.relu), new nn.Sequential(new nn.Linear(10, 10), new nn.ReLU()), new nn.Linear(10, 1), mx.sigmoid,);const y = model.forward(mx.random.normal([32, 2]));console.log(y);

目前暂时还没有 JavaScript API 的文档,开发者可以查阅可用 API 的 TypeScript 定义,以及 MLX 的官方网站以获取文档。

JavaScript API 基本上通过将 API 名称从 Snake_Case 转换为 CamelCase 来复制官方 Python API。 例如,mx.not_equal Python API 在 JavaScript 中重命名为 mx.notEqual。

由于 JavaScript 的限制,存在一些可能的例外情况:

JavaScript 数字始终是浮点值,因此 mx.array(42) 的默认数据类型是 mx.float32 而不是 mx.int32。mx.var API 已重命名为 mx.variance。运算符重载不起作用,可以使用 mx.add(a, b) 而不是 a + b。通过 [] 运算符进行索引不起作用,请使用 array.item 和 array.itemPut_ 方法代替删除数组不执行任何操作,必须等待垃圾回收才能释放数组的内存。Module 实例不能用作函数,必须使用 forward 方法

更多关于 node-mlx 的用法和知识可以参考文末资料,本文不再过多展开。

参考资料

https://github.com/ml-explore/mlx

https://github.com/frost-beta/node-mlx

https://www.youtube.com/watch?app=desktop&v=FplJsVd2dTk

https://medium.com/@kokodavidclaus/incorporating-machine-learning-models-into-node-js-backend-systems-448394bb08cf

0 阅读:0

前有科技后进阶

简介:感谢大家的关注