使用 MLX 和 Swift 进行设备端机器学习研究 [译]

Swift 编程语言在机器学习研究中展现出了巨大的潜力，因为它将 Python 语言的易用性与 C++ 语言的高效性完美结合。通过 Swift，研究人员可以在保持高效执行的同时，享受简洁的语法和易于理解的代码。

MLX 是专为 Apple Silicon 芯片设计的机器学习研究阵列框架。MLX 的设计目标是提供一个研究工具，而非用于应用程序中的模型生产部署。

MLX Swift 将 MLX 引入 Swift 语言，使得在 Apple Silicon 芯片上进行机器学习实验更加便捷。

此次发布包括以下内容：

• 完整的 Swift API，涵盖 MLX 核心功能
• 更高层次的神经网络和优化器包
• 使用 Mistral 7B 进行文本生成的示例
• MNIST 数据集训练的示例
• 提供 Swift 与 C++ 核心之间的桥梁：C API

所有内容都在宽松的 MIT 许可证下发布。

这标志着机器学习研究人员在 Swift 环境中进行实验的一个重要进展。

MLX 提供了一些现有 Swift 库几乎不具备的重要特性，包括：

• 原生硬件加速支持。MLX 可以在 CPU 或 GPU 上执行计算密集型任务。
• 自动微分功能，支持神经网络训练和基于梯度的机器学习模型。

欲了解更多 MLX 信息，请查看 文档。

Swift 语言以其快速、易用且与 Apple Silicon 芯片高度兼容而广受欢迎。借助 MLX Swift，研究人员能够在多个平台和设备上轻松进行机器学习实验，提升了跨平台研究的灵活性。

通过 Xcode 或 SwiftPM 设置 MLX Swift 非常简单且高效。

在 MLX Swift 中，N 维数组的构建与操作极为简便。以下示例展示了如何在默认设备（GPU，除非另有说明）上执行所有操作：

import MLX
import MLXRandom

let r = MLXRandom.normal([2])
print(r)
// array([-0.125875, 0.264235], dtype=float32)

let a = MLXArray(0 ..< 6, [3, 2])
print(a)
// array([[0, 1],
//        [2, 3],
//        [4, 5]], dtype=int32)

// 访问第 0 行的最后一个元素
print(a[0, -1])
// array(1, dtype=int32)

// 截取前两行
print(a[0 ..< 2])
// array([[0, 1],
//        [2, 3]], dtype=int32)

// 广播相加
let b = a + r

print(b)
// array([[-0.510713, 1.04633],
//        [1.48929, 3.04633],
//        [3.48929, 5.04633]], dtype=float32)

MLX Swift 中还支持函数变换。函数变换对于训练带有自动微分的模型以及优化计算图的速度或内存使用非常重要。以下示例展示了如何计算一个函数的梯度。

func fn(_ x: MLXArray) -> MLXArray {
    x.square()
}

let gradFn = grad(fn)

let x = MLXArray(1.5)
let dfdx = gradFn(x)

// 输出 2 * 1.5 = 3
print(dfdx)

文档中还提供了更多完整的 示例，帮助您快速入门 MLX Swift：

• 使用 LLM 生成文本：一个完整的 Mistral 7B 文本生成示例。此示例支持使用任何 Mistral 或 Llama 风格的模型生成文本，包括许多可在 Hugging Face 上找到的预量化 MLX 模型。
• 在 MNIST 数据集上训练 MLP：使用 MLX Swift 神经网络与优化器包训练一个简单的多层感知机，进行手写数字分类。

以下是一些其他有助于快速上手 MLX Swift 的资源：

• Swift 文档与示例
• GitHub 仓库
• 如果遇到问题或有改进建议，请提交 问题。
• 欢迎贡献！如果您有意参与 MLX Swift 的贡献，请查阅我们的 贡献指南。

原文链接：https://www.swift.org/blog/mlx-swift/