A side-by-side CNN implementation and comparability

Just a few months in the past, Apple quietly launched the primary public model of its MLX framework, which fills an area in between PyTorch, NumPy and Jax, however optimized for Apple Silicon. Very like these libraries, MLX is a Python-fronted API whose underlying operations are largely applied in C++.

Beneath are some observations of the similarities and variations between MLX and PyTorch. I applied a bespoke convolutional neural community utilizing PyTorch and its Apple Silicon GPU {hardware} assist, and examined it on a couple of completely different datasets. Particularly, the MNIST dataset, and the CIFAR-10 and CIFAR-100 datasets.

All of the code mentioned beneath could be discovered right here.

I applied the mannequin with PyTorch first, since I’m extra conversant in the framework. The mannequin has a collection of convolutional and pooling layers, adopted by a couple of linear layers with dropout.

# First block: Conv => ReLU => MaxPool
self.conv1 = Conv2d(in_channels=channels, out_channels=20, kernel_size=(5, 5), padding=2)
self.relu1 = ReLU()
self.maxpool1 = MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
# Second block: Conv => ReLU => MaxPool
self.conv2 = Conv2d(in_channels=20, out_channels=50, kernel_size=(5, 5), padding=2)
self.relu2 = ReLU()
self.maxpool2 = MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
# Third block: Conv => ReLU => MaxPool layers
self.conv3 = Conv2d(in_channels=50, out_channels=final_out_channels, kernel_size=(5, 5), padding=2)
self.relu3 = ReLU()
self.maxpool3 = MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
# Fourth block: Linear => Dropout => ReLU layers
self.linear1 = Linear(in_features=fully_connected_input_size, out_features=fully_connected_input_size // 2)
self.dropout1 = Dropout(p=0.3)
self.relu3 = ReLU()
# Fifth block: Linear => Dropout layers
self.linear2 = Linear(in_features=fully_connected_input_size // 2, out_features=fully_connected_input_size // 4)
self.dropout2 = Dropout(p=0.3)
# Sixth block: Linear => Dropout layers
self.linear3 = Linear(in_features=fully_connected_input_size // 4, out_features=lessons)
self.dropout3 = Dropout(p=0.3)
self.logSoftmax = LogSoftmax(dim=1)

This structure is overkill for MNIST dataset classification, however I needed one thing with some complexity to check the 2 frameworks. I examined this towards the CIFAR datasets, which approached round 40% accuracy; not superb, however I suppose respectable for one thing that isn’t a ResNet.

After ending this implementation, I wrote a parallel implementation leveraging MLX. I fortunately found that most of the PyTorch implementation might be instantly re-used, after importing the mandatory MLX modules and changing the PyTorch ones.

For instance, the MLX model of the above code is right here; it is similar other than a few variations in named parameters.

MLX has some fascinating properties value calling out.

Array

MLX’s array class takes the place of Tensor; a lot of the documentation compares it to NumPy’s ndarray, nevertheless it’s also the datatype used and returned by the assorted neural community layers obtainable within the framework.

array works principally as you’d count on, although I did have a little bit of hassle changing backwards and forwards between deeply-nested np.ndarrays and mlx.arrays necessitating some listing kind shuffling to make issues work.

Lazy Computation

Operations in MLX are lazily evaluated; that means that the one computation executed within the lazily-built compute graph is that which generates outputs truly used by this system.

There are two methods to drive analysis of the outcomes of operations (similar to inference):

• Calling mlx.eval() on the output
• Referencing the worth of a variable for any purpose; for instance when logging or inside conditional statements

This is usually a little difficult when attempting to handle the efficiency of the code, since a reference (even an incidental one) to any worth triggers an analysis of that variable in addition to all intermediate variables inside the graph. For instance:

def classify(X, y):
mannequin = MyModel()                   # Not but initialized
p = mannequin(X)                        # Not but computed
loss = mlx.nn.losses.nll_loss(p, y) # Not but computed
print(f"loss worth: {loss}") # Inits `mannequin`, computes `loss` _and_ `p`
mlx.eval(p)                  # No-op
# With out the print() above, would return `p` and lazy `loss`
return p, loss 

This conduct additionally makes slightly troublesome to construct one-to-one benchmarks between PyTorch and MLX-based fashions. Since coaching loops could not consider outputs inside the loop itself, its computation must be compelled in an effort to monitor the time of the particular operations.

test_start = time.perf_counter_ns() # Begin time block
accuracy, _ = eval(test_data_loader, mannequin, n)
mx.eval(accuracy) # Pressure calculation inside measurement block
test_end = time.perf_counter_ns() # Finish time block

There’s a tradeoff between accumulating a big implicit computation graph, and usually forcing the analysis of that graph throughout coaching. For instance, I used to be capable of lazily run via all of this mannequin’s coaching epochs over the dataset in only a few seconds. Nonetheless, the eventual analysis of that (presumably huge) implicit graph took roughly the identical period of time as eval’ing after every batch. That is in all probability not all the time the case.

Compilation

MLX offers the power to optimize the execution of pure capabilities via compilation. These could be both a direct name to mlx.compile() or an annotation (@mlx.compile) on a pure operate (with out uncomfortable side effects).

There are a couple of gotchas associated to state mutation when utilizing compiled capabilities; these are mentioned within the docs.

It looks like this ends in a compilation of logic into Steel Shader Language to be run on the GPU (I explored MSL earlier right here).

API Compatibility and Code Conventions

As talked about above, it was fairly simple to transform a lot of my PyTorch code into MLX-based equivalents. Just a few variations although:

• A number of the neural community layers discretely count on completely different configurations of inputs. For instance, mlx.nn.Conv2d expects enter photographs in NHWC format (with C representing the channels dimensionality), whereas torch.nn.Conv2d expects NCHW ; there are a couple of different examples of this. This required some conditional tensor/array shuffling.
• There may be sadly no analog to the relative pleasure which might be PyTorch Datasets and DataLoaders being at present supplied by MLX; as a substitute I needed to craft one thing resembling them by hand.
• Mannequin implementations, deriving from nn.Module, aren’t anticipated to override ahead() however fairly __call__() for inference
• I assume due to the potential for operate compilation, in addition to the lazy analysis assist talked about above, the method of coaching utilizing MLX optimizers is a bit completely different than with a typical PyTorch mannequin. Working with the latter, one is used to the usual format of one thing like:

for X, y in dataloader:
p = mannequin(X)
loss = loss_fn(p, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()

MLX encourages, and appears to count on, a format resembling the next, taken from the docs and one of many repository examples:

def loss_fn(mannequin, X, y):
return nn.losses.cross_entropy(mannequin(X), y, discount="imply")
loss_and_grad_fn = nn.value_and_grad(mannequin, loss_fn)
@partial(mx.compile, inputs=mannequin.state, outputs=mannequin.state)
def step(X, y):
loss, grads = loss_and_grad_fn(mannequin, X, y)
optimizer.replace(mannequin, grads)
return loss
# batch_iterate is a customized generator operate
for X, y in batch_iterate(batch_size, train_images, train_labels):
loss = step(X, y)

Which is okay, however a bit extra concerned than I used to be anticipating. In any other case, the whole lot felt very acquainted.

Notice that every one outcomes beneath are from my MacBook Air M2.

This CNN has three configurations: PyTorch CPU, PyTorch GPU, and MLX GPU. As a sanity test, over 30 epochs, right here’s how the three examine when it comes to accuracy and loss:

The outcomes listed here are all in the identical ballpark, although it’s fascinating that the MLX-based mannequin seems to converge extra shortly than the PyTorch-based ones.

As well as, it looks like the accuracy of the MLX mannequin is persistently barely beneath that of the PyTorch-based fashions. I’m undecided what accounts for that discrepancy.

When it comes to runtime efficiency, I had different fascinating outcomes:

When coaching the mannequin, the PyTorch-based mannequin on the CPU unsurprisingly took essentially the most time, from a minimal of 36 to a most of 45 seconds per epoch. The MLX-based mannequin, working on the GPU, had a variety of about 21–27 seconds per epoch. PyTorch working on the GPU, through the MPS machine , was the clear winner on this regard, with epochs starting from 10–14 seconds.

Classification over the check dataset of ten thousand photographs tells a special story.

Whereas it took the CPU-based mannequin round 1700ms to categorise all 10k photographs in batches of 512, the GPU-based fashions accomplished this job in 1100ms for MLX and 850ms for PyTorch.

Nonetheless, when classifying the photographs individually fairly than in batches:

Apple Silicon makes use of a unified reminiscence mannequin, which signifies that when setting the information and mannequin GPU machine to mps in PyTorch through one thing like .to(torch.machine(“mps”)) , there isn’t a precise motion of information to bodily GPU-specific reminiscence. So it looks like the overhead related to PyTorch’s initialization of Apple Silicon GPUs for code execution is pretty heavy. As seen additional above, it really works nice throughout parallel batch workloads. However for particular person file classification after coaching, it was far outperformed by no matter MLX is doing underneath the hood to spin up GPU execution extra shortly.

Profiling

Taking a fast have a look at some cProfile output for the MLX-based mannequin, ordered by cumulative execution time:

   ncalls  tottime  percall  cumtime  percall filename:lineno(operate)
426   86.564    0.203   86.564    0.203 {built-in technique mlx.core.eval}
1    2.732    2.732   86.271   86.271 /Customers/mike/code/cnn/src/python/mlx/cnn.py:48(prepare)
10051    0.085    0.000    0.625    0.000 /Customers/mike/code/cnn/src/python/mlx/mannequin.py:80(__call__)
30153    0.079    0.000    0.126    0.000 /Customers/mike/Library/Python/3.9/lib/python/site-packages/mlx/nn/layers/pooling.py:23(_sliding_windows)
30153    0.072    0.000    0.110    0.000 /Customers/mike/Library/Python/3.9/lib/python/site-packages/mlx/nn/layers/convolution.py:122(__call__)
1    0.062    0.062    0.062    0.062 {built-in technique _posixsubprocess.fork_exec}
40204    0.055    0.000    0.055    0.000 {built-in technique relu}
10051    0.054    0.000    0.054    0.000 {built-in technique mlx.core.imply}
424    0.050    0.000    0.054    0.000 {built-in technique step}

We a while spent right here in a couple of layer capabilities, with the majority of time spent in mlx.core.eval(), which is sensible because it’s at this level within the graph that issues are literally being computed.

Utilizing asitop to visualise the underlying timeseries powertools knowledge from MacOS:

You may see that the GPU is absolutely saturated through the coaching of this mannequin, at its most clock velocity of 1398 MHz.

Now examine to the PyTorch GPU variant:

   ncalls  tottime  percall  cumtime  percall filename:lineno(operate)
15585   41.385    0.003   41.385    0.003 {technique 'merchandise' of 'torch._C.TensorBase' objects}
20944    6.473    0.000    6.473    0.000 {built-in technique torch.stack}
31416    1.865    0.000    1.865    0.000 {built-in technique torch.conv2d}
41888    1.559    0.000    1.559    0.000 {built-in technique torch.relu}
31416    1.528    0.000    1.528    0.000 {built-in technique torch._C._nn.linear}
31416    1.322    0.000    1.322    0.000 {built-in technique torch.max_pool2d}
10472    1.064    0.000    1.064    0.000 {built-in technique torch._C._nn.nll_loss_nd}
31416    0.952    0.000    7.537    0.001 /Customers/mike/Library/Python/3.9/lib/python/site-packages/torch/utils/knowledge/_utils/collate.py:88(collate)
424    0.855    0.002    0.855    0.002 {technique 'run_backward' of 'torch._C._EngineBase' objects}
5    0.804    0.161   19.916    3.983 /Customers/mike/code/cnn/src/python/pytorch/cnn.py:176(eval)

Curiously, the highest operate seems to be Tensor.merchandise(), which is named in varied locations within the code to calculate loss and accuracy, and presumably additionally inside a number of the layers referenced decrease within the stack. Eradicating the monitoring of loss and accuracy throughout coaching would in all probability have a noticeable enchancment on general coaching efficiency.

In comparison with the MLX mannequin, the PyTorch variant doesn’t appear to have saturated the GPU throughout coaching (I didn’t see it breach 95%), and has the next stability of utilization on the CPU’s E cores and P cores.

It’s fascinating that the MLX mannequin makes heavier use of the GPU, however trains significantly extra slowly.

Neither mannequin (CPU or GPU-based) seems to have engaged the ANE (Apple Neural Engine).

MLX was simple to choose up, and that must be the case for anybody with expertise utilizing PyTorch and NumPy. Although a number of the developer documentation is a bit skinny, given the intent to offer instruments suitable with these frameworks’ APIs, it’s simple sufficient to fill in any gaps with the corresponding PyTorch or NumPy docs (for instance, SGD [1] [2]).

The general efficiency of the MLX mannequin was fairly good; I wasn’t positive whether or not I used to be anticipating it to persistently outperform PyTorch’s mps machine assist, or not. Whereas it appeared like coaching was significantly sooner via PyTorch on the GPU, single-item prediction, significantly at scale, was a lot sooner via MLX for this mannequin. Whether or not that’s an impact of of my MLX configuration, or simply the properties of the framework, its exhausting to say (and if its the previous — be at liberty to depart a difficulty on GitHub!)