Mac-optimized TensorFlow flexes new M1 and GPU muscles

A new Mac-optimized fork of machine learning environment TensorFlow posts some major performance increases. Although a big part of that is that until now the GPU wasn’t used for training tasks (!), M1-based devices see even further gains, suggesting a spate of popular workflow optimizations like this one are incoming.

Announced on both TensorFlow and Apple’s blogs, the improved Mac version shows in the best case more than a 10x improvement in speed for common training tasks.

That’s worth celebrating on its own for anyone who works in ML and finds themselves constantly waiting for their models to bake. But the fact that previous versions of TF only utilized the CPU on Macs and not the powerful parallel processors in the GPU probably limited the pool of people who inflict that problem on themselves in the first place. (Most large-scale ML training is done using cloud computing.)

The change from CPU-only to CPU+GPU could account for a great deal of the improvement, as the benchmarks on an Intel-based Mac Pro show huge gains on the same hardware. Training times once in the 6-8 second range are now measured in fractions of a second.

That’s not to say the M1 isn’t capable, but the new M1 Macs also have new GPUs, meaning the jump from nearly 10 seconds for a task on a 2019 MacBook Pro to less than 2 on a new M1 machine can only be partly attributed to Apple’s fancy first-party silicon.

I contacted Apple for more information, such as a number for an M1 device running non-optimized code (which would elucidate the improvements nicely) but a representative said it does not have those numbers.

At any rate perhaps more important for developers will be the improved battery life and heat management of the M1 devices. Performance bumps are all well and good, but if it made your machine into a hot plate, blasted your fan and made you run for the outlet in under an hour — not so good. Fortunately the M1 seems to be demonstrating remarkable efficiency under load, neither draining its reserves or heating up too much.

You can probably expect a lot of these “now works better on M1” stories now that the new Macs are out and all the major companies can ship the updates they’ve been sitting on for the last few months.