MLX: 애플 전용칩을 위한 머신러닝 프레임워크 / MLX: An array framework for Apple silicon
소개
Apple
의 ML팀에서 Apple Silicon용 머신러닝 프레임워크인 MLX를 공개하였습니다. 그것도 무려 MIT 라이선스로 공개하였습니다. 
공식 소개에는 MLX를 Array Framework라고 하지만, 문서를 보면 ReLU, Conv, MultiHeadAttention, Embedding 등과 같은 주요 Neural Network의 Layer들은 물론, 최적화를 위한 Optimizer나 Loss 등을 제공하고 있어, 그냥 ML/DL framework라고 불러도 될 것 같습니다. 
아래는 공식 문서의 MLP(Multi-Layer Perceptron) 예시인데요, 보시면 NumPy/PyTorch스럽게 모델을 구성하고 학습/추론하는 것을 보실 수 있습니다.
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np
class MLP(nn.Module):
def __init__(
self, num_layers: int, input_dim: int, hidden_dim: int, output_dim: int
):
super().__init__()
layer_sizes = [input_dim] + [hidden_dim] * num_layers + [output_dim]
self.layers = [
nn.Linear(idim, odim)
for idim, odim in zip(layer_sizes[:-1], layer_sizes[1:])
]
def __call__(self, x):
for l in self.layers[:-1]:
x = mx.maximum(l(x), 0.0)
return self.layers[-1](x)
def loss_fn(model, X, y):
return mx.mean(nn.losses.cross_entropy(model(X), y))
def eval_fn(model, X, y):
return mx.mean(mx.argmax(model(X), axis=1) == y)
num_layers = 2
hidden_dim = 32
num_classes = 10
batch_size = 256
num_epochs = 10
learning_rate = 1e-1
# Load the data
import mnist
train_images, train_labels, test_images, test_labels = map(
mx.array, mnist.mnist()
)
def batch_iterate(batch_size, X, y):
perm = mx.array(np.random.permutation(y.size))
for s in range(0, y.size, batch_size):
ids = perm[s : s + batch_size]
yield X[ids], y[ids]
# Load the model
model = MLP(num_layers, train_images.shape[-1], hidden_dim, num_classes)
mx.eval(model.parameters())
# Get a function which gives the loss and gradient of the
# loss with respect to the model's trainable parameters
loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
# Instantiate the optimizer
optimizer = optim.SGD(learning_rate=learning_rate)
for e in range(num_epochs):
for X, y in batch_iterate(batch_size, train_images, train_labels):
loss, grads = loss_and_grad_fn(model, X, y)
# Update the optimizer state and model parameters
# in a single call
optimizer.update(model, grads)
# Force a graph evaluation
mx.eval(model.parameters(), optimizer.state)
accuracy = eval_fn(model, test_images, test_labels)
print(f"Epoch {e}: Test accuracy {accuracy.item():.3f}")
주요 특징
MLX는 머신러닝 연구자들이, 머신러닝 연구자들을 위해 설계 및 개발하였습니다. MLX는 사용자 친화적이면서도 효율적으로 모델을 학습 및 배포할 수 있도록 고안되었습니다.
MLX의 설계는 NumPy, PyTorch, Jax, ArrayFire와 같은 프레임워크에서 영감을 받았으며, 연구자들이 MLX를 사용하여 새로운 아이디어를 빠르게 탐색할 수 있도록 하는 것이 목표입니다.
친숙한 API / Familiar APIs
NumPy와 매우 유사한 Python API와 다양한 기능의 C++ API를 제공합니다. 또한, PyTorch의 torch.nn이나, torch.optim과 같은 mlx.nn 및 mlx.optimizers를 제공하여 복잡한 모델을 간단히 구성할 수 있도록 합니다.
다양한 연산 함수 제공 / Composable function transformations
자동 미분(Automatic Differentiation), 자동 벡터화(Automatic Vectorization), 연산 그래프 최적화(Computation Graph Optimization)와 같은 기능들(Composable function transformations)을 제공합니다.
지연 연산 / Lazy computation
Lazy computation을 적용하여 필요할 때까지 연산 및 배열의 구체화를 늦춥니다.
동적 그래프 구성 / Dynamic graph construction
연산 그래프를 동적으로 구성하여 함수 인자의 모양을 변경하더라도 컴파일이 느려지지 않으며, 간단하고 직관적으로 디버깅할 수 있습니다.
CPU 및 GPU 지원 / Multi-device
현재는 CPU와 GPU 모두에서 연산을 실행할 수 있습니다.
통합 메모리 / Unified Memory
공유 메모리(shared memory)에 배열(array)을 저장하여 데이터를 이동하지 않고도 지원하는 장치들에서의 연산이 가능합니다. 이는 다른 프레임워크와 눈에 띄는 차별점입니다.
MLX 공식 예제들
더 읽어보기
딥러닝 프레임워크 MLX의 GitHub 저장소
MLX 공식 문서
https://ml-explore.github.io/mlx/build/html/index.html