Машинное обучение на устройствах с помощью MLX и Swift

Язык программирования Swift обладает большим потенциалом для использования в работе с машинным обучением, поскольку сочетает в себе простоту использования и высокоуровневый синтаксис такого языка, как Python, и скорость компилируемого языка, такого как C++.

MLX — это фреймворк для работы с машинным обучением на Apple Silicon. MLX предназначен для исследований, а не для развертывания моделей в приложениях.

MLX Swift расширяет MLX на язык Swift, облегчая исследователям ML-эксперименты на процессорах Apple.

В этот релизе включены:

• Всеобъемлющий Swift API для ядра MLX
• Пакеты нейронных сетей и оптимизаторов высокого уровня
• Пример генерации текста с помощью Mistral 7B
• Пример обучения MNIST
• C API для MLX, который служит мостом между Swift и ядром C++

Все вышеперечисленное выпущено под разрешительной лицензией MIT.

Это большой шаг к тому, чтобы ML-исследователи могли экспериментировать с ИИ, используя Swift.

Мотивация

MLX обладает несколькими важными для исследований в области машинного обучения функциями, которые поддерживают лишь немногие из существующих библиотек Swift. К ним относятся:

• Нативная поддержка аппаратного ускорения. MLX может выполнять операции с интенсивными вычислениями на CPU или GPU.
• Автоматическая дифференциация для обучения нейронных сетей и моделей машинного обучения на основе градиента.

Более подробную информацию о MLX можно найти в документации.

Язык программирования Swift — быстрый, простой в использовании и хорошо работает на Apple Silicon. С помощью MLX Swift вы получаете удобный для исследователей фреймворк машинного обучения с возможностью легко проводить эксперименты на различных платформах и устройствах.

Обзор

Начать работу с MLX Swift можно быстро и легко с помощью Xcode или SwiftPM.

В 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)

// last element of 0th row
print(a[0, -1])
// array(2, dtype=int32)

// slice of the first two rows        
print(a[0 ..< 2])
// array([[0, 1],
//        [2, 3]], dtype=int32)

// add with broadcast
let b = a + r

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

Вы также можете использовать преобразования функций в MLX Swift. Преобразования функций в MLX полезны для обучения моделей с автоматическим дифференцированием, а также для оптимизации вычислительных графов по скорости или использованию памяти. Ниже приведен пример, который вычисляет градиент функции.

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

let gradFn = grad(fn)

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

// prints 2 * 1.5 = 3
print(dfdx)

В документации есть несколько более полных примеров, которые помогут вам начать работу с MLX Swift:

• Генерация текста с помощью LLM: Полный пример генерации текста LLM с помощью Mistral 7B. Пример генерирует текст, используя любую модель в стиле Mistral или Llama, включая предварительно квантованные модели MLX, многие из которых доступны на Hugging Face.
• Обучение MLP на MNIST: Пример обучает простой многослойный перцептрон для классификации цифр MNIST, используя пакеты нейронных сетей и оптимизаторов MLX Swift.

Дополнительные ресурсы

Вот еще несколько ресурсов для начала работы с MLX Swift:

• Документация и примеры на Swift
• Репозиторий GitHub
• Если у вас возникли проблемы или есть предложения по улучшению, рекомендуется создать issue
• Если вы хотите внести свой вклад в MLX Swift, ознакомьтесь с руководством по внесению вклада.

Источник