Apple开源Apple Silicon机器学习框架MLX
时间:2023-12-06 16:41:05 小新
12月6日 消息:Apple 开源的 MLX 是一个适用于苹果芯片的机器学习框架。它具有许多功能,包括熟悉的 API、可组合的函数转换、延迟计算、动态图构建和多
12月6日 消息:Apple 开源的 MLX 是一个适用于苹果芯片的机器学习框架。它具有许多功能,包括熟悉的 API、可组合的函数转换、延迟计算、动态图构建和多设备支持等。
MLX 从 NumPy、PyTorch、Jax 和 ArrayFire 等框架中吸取灵感,创建了一个强大且多功能的平台。通过示例仓库,我们可以看到 MLX 的各种能力,包括语言模型训练、文本生成、图像生成和语音识别。MLX 的安装也非常简单,只需使用 pip 命令即可。
代码: https://github.com/ml-explore/mlx
文档: https://ml-explore.github.io/mlx/build/html/index.html
重点功能:
熟悉的 API:MLX 的 Python API 与 NumPy 紧密结合,而功能齐全的 C++ API 则反映了 Python 版本。此外,mlx.nn 和 mlx.optimizers 等更高级别的包通过遵守 PyTorch 约定来简化模型构建。
可组合函数转换:MLX 引入了可组合函数转换,支持自动微分、向量化和计算图优化。
惰性计算:MLX 中的计算被设计为惰性计算,确保仅在必要时才具体化数组,从而优化计算效率。
动态图构建:MLX采用动态图构建,消除了因函数参数形状变化而引发的缓慢编译。这种方法简化了调试过程。
多设备支持:MLX 允许操作在支持的设备(包括 CPU 和 GPU)上无缝运行,为开发人员提供了灵活性。
统一内存模型:MLX 引入了与其他框架不同的统一内存模型。阵列驻留在共享内存中,允许跨不同设备类型对 MLX 阵列进行操作,而无需移动数据。