[FIXED] How to move PyTorch model to GPU on Apple M1 chips?

Issue

On 18th May 2022, PyTorch announced support for GPU-accelerated PyTorch training on Mac.

I followed the following process to set up PyTorch on my Macbook Air M1 (using miniconda).

conda create -n torch-nightly python=3.8 

$ conda activate torch-nightly

$ pip install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

I am trying to execute a script from Udacity's Deep Learning Course available here.

The script moves the models to GPU using the following code:

G.cuda()
D.cuda()

However, this will not work on M1 chips, since there is no CUDA.

If we want to move models to M1 GPU and our tensors to M1 GPU, and train entirely on M1 GPU, what should we be doing?

---

If Relevant: G and D are Discriminator and Generators for GAN's.

class Discriminator(nn.Module):

    def __init__(self, conv_dim=32):
        super(Discriminator, self).__init__()
        self.conv_dim = conv_dim
        # complete init function
        self.cv1 = conv(in_channels=3, out_channels=conv_dim, kernel_size=4, stride=2, padding=1, batch_norm=False)            # 32*32*3  -> 16*16*32
        self.cv2 = conv(in_channels=conv_dim, out_channels=conv_dim*2, kernel_size=4, stride=2, padding=1, batch_norm=True)    # 16*16*32 -> 8*8*64
        self.cv3 = conv(in_channels=conv_dim*2, out_channels=conv_dim*4, kernel_size=4, stride=2, padding=1, batch_norm=True)  # 8*8*64   -> 4*4*128
        self.fc1 = nn.Linear(in_features = 4*4*conv_dim*4, out_features = 1, bias=True)
        

    def forward(self, x):
        # complete forward function
        out = F.leaky_relu(self.cv1(x), 0.2)
        out = F.leaky_relu(self.cv2(x), 0.2)
        out = F.leaky_relu(self.cv3(x), 0.2)
        out = out.view(-1, 4*4*conv_dim*4)
        out = self.fc1(out)
        return out    

D = Discriminator(conv_dim)

class Generator(nn.Module):    
    def __init__(self, z_size, conv_dim=32):
        super(Generator, self).__init__()
        self.conv_dim = conv_dim
        self.z_size = z_size
        # complete init function
        self.fc1 = nn.Linear(in_features = z_size, out_features = 4*4*conv_dim*4)
        self.dc1 = deconv(in_channels = conv_dim*4, out_channels = conv_dim*2, kernel_size=4, stride=2, padding=1, batch_norm=True)
        self.dc2 = deconv(in_channels = conv_dim*2, out_channels = conv_dim, kernel_size=4, stride=2, padding=1, batch_norm=True)
        self.dc3 = deconv(in_channels = conv_dim, out_channels = 3, kernel_size=4, stride=2, padding=1, batch_norm=False)

    def forward(self, x):
        # complete forward function
        x = self.fc1(x)
        x = x.view(-1, conv_dim*4, 4, 4)
        x = F.relu(self.dc1(x))
        x = F.relu(self.dc2(x))
        x = F.tanh(self.dc3(x))
        return x

G = Generator(z_size=z_size, conv_dim=conv_dim)

Solution

This is what I used:

if torch.backends.mps.is_available():
    mps_device = torch.device("mps")
    G.to(mps_device)
    D.to(mps_device)

Similarly for all tensors that I want to move to M1 GPU, I used:

tensor_ = tensor_(mps_device)

Some operations are ot yet implemented using MPS, and we might need to set a few environment variables to use CPU fall back instead: One error that I faced during executing the script was

# NotImplementedError: The operator 'aten::_slow_conv2d_forward' is not current implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.

To solve it I set the environment variable PYTORCH_ENABLE_MPS_FALLBACK=1

conda env config vars set PYTORCH_ENABLE_MPS_FALLBACK=1
conda activate <test-env>

References:

1. https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/
2. https://pytorch.org/docs/master/notes/mps.html
3. https://sebastianraschka.com/blog/2022/pytorch-m1-gpu.html
4. https://sebastianraschka.com/blog/2022/pytorch-m1-gpu.html
5. https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#setting-environment-variables

Answered By - Anirban Saha