M1 macbook using GPU for training: problem and solution

[yingjoeyzhou]

Short description of the problem

Having set up the osld environment following steps below (see screenshot), I was able to train HMMs on my M1 macbook. However, tensorflow did not seem to use GPU at all.

Solution

As suggested by https://developer.apple.com/metal/tensorflow-plugin/ , seems like installing the tensorflow-metal plugin solves the problem. Concretely, after pip install tensorflow-macos==2.9.1, I also ran pip install tensorflow-metal==0.5.0. This works in my case as in, if I evaluate

import tensorflow as tf
tf.config.list_physical_devices('GPU')

I now get [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')].