在 Apple Silicon 芯片的 Mac 上安装深度学习环境

在 Apple Silicon 芯片的 Mac 上安装深度学习环境的最小示例。

说明：

1. 苹果提高了安装 tensorflow 的相关依赖以及安装教程[^1]

2. pytorch 1.12 开始支持苹果自研芯片携带的 GPU，设备标记为 Metal Performance Shaders (mps)[^2]

conda create --name your_env  python=3.9
conda activate your_env

# install tensorflow
# also see https://developer.apple.com/metal/tensorflow-plugin/
conda install -c apple tensorflow-deps
python -m pip install tensorflow-macos
python -m pip install tensorflow-metal

# install pytorch
# follow official instruction
conda install -c pytorch pytorch torchvision torchaudio

# install jax and flax
conda install -c conda-forge jax jaxlib flax

# install transformers
# pip is better than conda
pip install transformers datasets

# install scientific packages
conda install pandas numpy scipy scikit-learn matplotlib

# install other packages
conda install jupyter jsonlines

安装完成，检查 tf 安装情况。可以看到苹果的 GPU！

In [1]: import tensorflow as tf

In [2]: tf.config.list_physical_devices()
Out[2]: 
[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'),
 PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

检查 pytorch 的安装情况。可以看到 mps 是可用的，后续使用将“cuda”替换为“mps”即可。

In [1]: import torch


In [3]: torch.backends.mps.is_available()
Out[3]: True

你还可以使用 https://github.com/rasbt/machine-learning-notes/benchmark/pytorch-m1-gpu 提供的脚本来测试 mps 和 cpu 版本的程序性能。

[^1]: Tensorflow Plugin - Metal - Apple Developer：https://developer.apple.com/metal/tensorflow-plugin/

[^2]: PyTorch 1.12: TorchArrow, Functional API for Modules and nvFuser, are now available | PyTorch：https://pytorch.org/blog/pytorch-1.12-released/