在Mac M1/M2上安装支持GPU的PyTorch和TensorFlow

mac m1,m2 安装 提供GPU支持的pytorch和tensorflow

mac m1 刚出的时候，各种支持都不完善。那时候要使用conda，只能选择miniconda。几年过去了，各种主流软件对mac m1，m2的支持都已经非常完善了。

比如Pytorch，正如官网所写：

In collaboration with the Metal engineering team at Apple, we are excited to announce support for GPU-accelerated PyTorch training on Mac.

于是，我重新基于Anaconda安装了一下GPU支持的Pytorch和Tensorflow。

Anaconda

安装

Anaconda 下载地址：Anaconda 下载
可以看到官方已经支持M1 Chip。

下载图形化安装包或者命令行安装方式到本地。我更喜欢命令行的安装。
然后执行即可开始安装：

bash 安装白路径/安装包名
# 或者
zsh 安装白路径/安装包名

测试

测试是否安装成功：

conda --version

如果成功输出，则代表安装成功。（安装完成之后要source一下~/.zshrc或者重启终端）

Pytorch

参考链接

安装步骤

安装 Xcode

安装过的可以直接到下一步。
命令行安装：xcode-select --install，或者App Store下载。

创建conda环境

conda create -n torch-gpu python=3.8
conda activate torch-gpu

然后到Pytorch官网选择对应的版本，复制安装命令。

测试

终端输入python进入python命令行。

import torch
import math

print(torch.backends.mps.is_available())
# True
print(torch.backends.mps.is_built())
# True

都输出为True即可。

加速效果

注意

To run PyTorch code on the GPU, use torch.device(“mps”) analogous to torch.device(“cuda”) on an Nvidia GPU.

在mac m1上，device是’mps’ 而不是’cuda’。

Tensorflow

参考链接

安装步骤

安装 Xcode

安装过的可以直接到下一步。
命令行安装：xcode-select --install，或者App Store下载。

指定安装环境

name: tf-metal
channels:
  - apple
  - conda-forge
dependencies:
  - python=3.9  ## specify desired version
  - pip
  - tensorflow-deps

  ## uncomment for use with Jupyter
  ## - ipykernel

  ## PyPI packages
  - pip:
    - tensorflow-macos
    - tensorflow-metal  ## optional, but recommended

然后使用conda，根据文件内容创建新环境。

conda env create -n tf_gpu -f tf-metal-arm64.yaml

加速效果

测试

import sys
import tensorflow.keras
import tensorflow as tf
import platform
print(f"Python Platform: {platform.platform()}")
print(f"Tensor Flow Version: {tf.__version__}")
print(f"Keras Version: {tensorflow.keras.__version__}")
print()
print(f"Python {sys.version}")

gpu = len(tf.config.list_physical_devices('GPU'))>0
print("GPU is", "available" if gpu else "NOT AVAILABLE")

The End

安装完成。