MacOS M2 Tensorflow GPU环境搭建

本文记录在mac mini m2 pro上搭建tensorflow gpu版本的过程,中间因为一些坑,导致环境搭建过程有点反复,希望对有相同需求的同学有帮助。

1. 安装conda环境

这里创建conda环境的时候需要指定python版本为3.9,这是因为后面需要安装的tensorflow需要python 3.6~3.9

1
conda create -n tf_gpu python=3.9

2. 安装tensorflow

安装tensorflow的时候,需要指定版本为2.14.1,这里是因为后面需要安装的tensorflow-metal,最新版本只支持2.14

1
pip install tensorflow==2.14.1

3. 安装tensorflow-metal

目前tensorflow官方没有支持apple gpu,需要通过tensorflow-metal插件来支持使用apple gpu

1
pip install tensorflow-metal==1.1.0

4. 测试验证

1
2
3
4
5
import tensorflow.keras
import tensorflow as tf
print(f"Tensor Flow Version: {tf.__version__}")
gpu = len(tf.config.list_physical_devices('GPU'))>0
print("GPU is", "available" if gpu else "NOT AVAILABLE")

成功后会得到下面的输出

1
2
3
4
Tensor Flow Version: 2.14.1
Scikit-Learn 1.5.0
SciPy 1.13.1
GPU is available