MacBook 上安装 Tensorflow 支持 MPS GPU

MacBook 上安装测试 Tensorflow，Pytorch 对MPS GPU的支持

作者平时使用苹果的 MacBook M3 笔记本，在学习机器学习过程经常因为没有Nvidia显卡和CUDA环境导致效率低。MacBook M3 机器上带有强大的显卡 MPS (Metal Performance Shaders)，并且 Tensorflow 和 Pytorch 都支持基于 MPS 的 GPU 计算。下面记录一下在MacOS环境中安装支持 MPS 的基本步骤。

Tensorflow 环境

1. 创建运行环境

为了保证各个项目之间的开发环境隔离，最好每个项目单独创建 python 得运行环境。我平时使用的是mini conda来管理 python 运行环境。

# 创建一个名为 tf_mps 的 python 运行环境，python 版本为3.10
% conda create -n tf_mps python=3.10

# 激活 tf_mps 环境
% conda activate tf_mps

2. 安装开发库

注意，安装默认的 tensorflow 开发库后，是不支持对 Metal GPU 的调用。必须还要安装 tensorflow-metal 库。

# 安装 tensorflow 开发库
% pip install tensorflow
# 安装 tensorflow 环境下对 Metal GPU 的支持库
% pip install tensorflow-metal

3. 测试是否支持GPU

安装完成后，测试一下对 Metal GPU 的支持是否启用。注意，如果是在 jupyter-lab 或是 VScode 的 jupyter-lab 环境中，安装完 tensorflow-metal 后，需要重启一下环境。

然后编写如下 python 脚本测试，测试结果

import tensorflow as tf

# 打印 tensorflow的版本，需要大于V2.9版本才支持Metal GPU
print(tf.__version__)

# 打印当前机器上的物理计算设备
print(tf.config.list_physical_devices())

输入结果示例如下：

2.17.0
[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'),
 PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

可见当前 tensorflow 的版本是 2.17.0，计算设备包括一个 CPU 和 一个 GPU 单元。说明当前的 tensorflow 已经支持 Metal GPU 的计算。

Pytorch 环境

1. 安装

pytorch 的安装比较简单，创建3.8或更高的python环境。直接到 pytorch官网 选择Mac环境，然后复制安装命令执行即可。

pip3 install torch torchvision torchaudio

2. 测试

写段代码简单测试一下是否支持mps

import torch
if torch.backends.mps.is_available():
    mps_device = torch.device("mps")
    x = torch.ones(1, device=mps_device)
    print (x)
else:
    print ("MPS device not found.")