Training Deep Learning Models on M1 Macbooks

Requirements: M1 Macbook with PyTorch Installed

Up until recently, it was both expensive and difficult to train PyTorch models on Macbooks because most Macbooks did not have cuda GPUs. However, with the introduction of Apple's own CPU/GPU chips, I was excited at the prospect of Macbooks finally being supported in PyTorch. This blog post covers how to set up PyTorch so that you can train on M1 Macbooks. To get started, install the nightly build of PyTorch using the following command. Note that this command is likely to change in the future as this will likely be available in the next stable release.

conda install pytorch torchvision torchaudio -c pytorch-nightly

Checking Compatability

Next, we can check if the current machine meets system requuirements and the correct and compatible PyTorch version was installed.

import torch

if torch.backend.mps.is_available() && torch.backend.mps.is_built():
    print("You have correct hardware and the proper pytorch version installed!")

Setting the device

Now that we have verified the pytorch version and MacOS version, we need to specify what device we want to train our models on. In this case, we want it to be on our local Macbook GPU

device = torch.device('mps')

# Set model device and other tensors' device so that it will be trained appropriately
model.to(device) 

Closing

With this simple change, you can now start training your models on your Macbook M1 machines. In the future, I would like to compare the difference in performance for the CPU and GPU of the M1 Macbook Pro when it comes to training state-of-the-art pretrained models.