Apple Silicon M1 でtensorflow-macosを実行したらめちゃくちゃ速かった。

はじめに

Macbook Air (AppleSilicon, M1)を購入しました。
Appleが設計したM1は下馬評以上の性能を叩き出し、とても盛り上がっていますね。

M1のハードウェアとしての魅力はもちろんすごいですが、M1に合わせた各種ソフトウェアの最適化も魅力的です。

AppleがTensorflowをフォークしてM1で最高のパフォーマンスを発揮するように最適化したコード(tensorflow-macos)を公開しています。
https://github.com/apple/tensorflow_macos#requirements

tensorflowのブログ記事でパフォーマンスが比較されていますが、前世代のIntel Macと比べると信じられないほどの性能向上を見ることができます。
https://blog.tensorflow.org/2020/11/accelerating-tensorflow-performance-on-mac.html

この記事では

1. 実際にM1 Macでtensorflow-macosを動作させた過程を紹介します。
2. サンプルコードを動作させてM1 Mac, Intel Mac, RTX 2080との実行時間の比較をします。

目次

1. tensorflow-macosを実行する環境を構築する。
2. 実行時間を比較する。
3. 考察、議論
4. 参考文献

tensorflow-macosを実行する環境を構築する。

Python 3.8 (for ARM)をインストールする。

tensorflow-macos 0.1a0はPython 3.8で動作します。
tensorflow-macosリポジトリのREQUIREMENTSにあるように、
必ずXcode Command Line Toolsからインストールしてください。

$ xcode-select --install

$ python3 -V
Python 3.8.2

$ which python3 
/usr/bin/python3

tensorflow_macosはARMアーキテクチャ向けのコードになっているので、PythonもARMアーキテクチャ向けのものにする必要があります。

例えば以下のようにインストールしたものはtensorflow-macosの実行時にwrong architectureというようなエラーが発生しました。

・ Python公式サイトからdmgファイルをダウンロードしてインストール
・ Anacondaを使ってインストール
・ pyenvを使ってインストール

Traceback (most recent call last):
  File "/path/to/tensorflow_macos_venv/lib/python3.8/site-packages/tensorflow/python/pywrap_tensorflow.py", line 64, in <module>
    from tensorflow.python._pywrap_tensorflow_internal import *
ImportError: dlopen(/path/to/tensorflow_macos_venv/lib/python3.8/site-packages/tensorflow/python/_pywrap_tensorflow_internal.so, 6): no suitable image found.  Did find:
    /path/to/tensorflow_macos_venv/lib/python3.8/site-packages/tensorflow/python/_pywrap_tensorflow_internal.so: mach-o, but wrong architecture
    /path/to/tensorflow_macos_venv/lib/python3.8/site-packages/tensorflow/python/_pywrap_tensorflow_internal.so: mach-o, but wrong architecture

tensorflow-macosの仮想環境(venv)を構築する。

tensorflow_macosの実行に必要な仮想環境を構築してくれる便利なスクリプトをAppleが用意してくれています。

tensorflow-macosリポジトリのINSTALLATIONにあるように、ターミナルからスクリプトを実行します。

$ /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/apple/tensorflow_macos/master/scripts/download_and_install.sh)"

仮想環境のパッケージのダウンロードが始まるので、しばらく待ちます(結構待つ)。
仮想環境を構築するパスを聞かれるので、Enterキーを押してデフォルトパスを指定し、構築を開始します。

Downloading installer.
/var/folders/hn/g5s5s0fx6zs2plvygdc95y140000gn/T/tmp.haPEac7f ~/Projects/tensorflow-mac
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100   659  100   659    0     0   1331      0 --:--:-- --:--:-- --:--:--  1331
100  316M  100  316M    0     0  2467k      0  0:02:11  0:02:11 --:--:-- 1638k
Extracting installer.
Path to new or existing virtual environment [default: /Users/tomoyaeibu/tensorflow_macos_venv/]: 
Using python from /usr/bin/python3.

###########################################################################

./install_venv.sh will perform the following operations:
-> Create new python 3.8 virtual environment at /Users/tomoyaeibu/tensorflow_macos_venv.
-> Install tensorflow_macos 0.1a0 into created virtual environment /Users/tomoyaeibu/tensorflow_macos_venv.
-> Install bundled binary wheels for numpy-1.18.5, grpcio-1.33.2, h5py-2.10.0, tensorflow_addons-0.11.2+mlcompute into /Users/tomoyaeibu/tensorflow_macos_venv.


Confirm [y/N]? 

tensorflow-macosの仮想環境(venv)をアクティベートする。

これでtensorflow-macosを実行する環境は整いました。
簡単ですね・・・！

tensorflow-macosの仮想環境をアクティベートします。

$ source /Users/{username}/tensorflow_macos_venv/bin/activate

$ pip list         
Package                Version
---------------------- ---------
absl-py                0.11.0
appnope                0.1.0
astunparse             1.6.3
backcall               0.2.0
cached-property        1.5.2
cachetools             4.1.1
certifi                2020.11.8
chardet                3.0.4
decorator              4.4.2
flatbuffers            1.12
gast                   0.4.0
google-auth            1.23.0
google-auth-oauthlib   0.4.2
google-pasta           0.2.0
grpcio                 1.33.2
h5py                   2.10.0
idna                   2.10
ipython                7.19.0
ipython-genutils       0.2.0
jedi                   0.17.2
Keras-Preprocessing    1.1.2
Markdown               3.3.3
numpy                  1.18.5
oauthlib               3.1.0
opt-einsum             3.3.0
parso                  0.7.1
pexpect                4.8.0
pickleshare            0.7.5
pip                    20.2.4
prompt-toolkit         3.0.8
protobuf               3.14.0
ptyprocess             0.6.0
pyasn1                 0.4.8
pyasn1-modules         0.2.8
Pygments               2.7.2
requests               2.25.0
requests-oauthlib      1.3.0
rsa                    4.6
setuptools             50.3.2
six                    1.15.0
tensorboard            2.4.0
tensorboard-plugin-wit 1.7.0
tensorflow-addons      0.11.2
tensorflow-estimator   2.3.0
tensorflow-macos       0.1a0
termcolor              1.1.0
traitlets              5.0.5
typeguard              2.10.0
typing-extensions      3.7.4.3
urllib3                1.26.2
wcwidth                0.2.5
Werkzeug               1.0.1
wheel                  0.35.1
wrapt                  1.12.1

tensorflow-macosがインストールされていることがわかります。

tensorflow-macosを実行する。

tensorflow-macosを実行するために既存のtensorflowを使ったトレーニングやバリデーションのコードを修正する必要はありません。

今まで使っていたスクリプトをそのまま使えます。
（ただしKerasを使っていた場合、ARM版のKerasがまだないようなので、tensorflowのAPIに置き換えるか、tf.KerasのAPIに置き換えてください。）

追記

@asamiKAさんがApple Silicon M1 は自然言語処理も、ちょっと速いよにて、より実践的な問題に対してtensorflow-macを適用しています。

記事によると、tensorflow-macはtensorflowで使うことのできるすべてのレイヤーについて最適化しているわけではないようです。
最適化されていないレイヤーを使う場合は、CPUのみが使われるため、学習スピードが遅くなってしまうはずです。
記事では最適化されていないGRUレイヤーを、最適化されているLSTMレイヤーに置き換えています。

将来的にtensorflow-macはtensorflowにマージされる予定があるらしいのですが、その時までには対応するレイヤーが増えていくのだと思います。

実行時間を比較する。

早速サンプルコードを実行して実行時間を比較してみましょう。
トレーニングにかかる時間を比較します。

実行環境

以下の3つの環境で比較しました。

	Macbook Air (2020,M1)	Macbook Pro (2017,Intel)	開発に使っている学習用PC
OS	Mac Big Sur	Mac Big Sur	Ubuntu 20.04.1
CPU	M1	2.3GHzデュアルコアIntel Core i5	Core i7-8700 3.2GHz
GPU	M1	Intel Iris Plus Graphics 640	RTX 2080 8GB
RAM	8GB	8GB	16GB
Python	python 3.8.2	python 3.8.2	python 3.8.2
Tensorflow	tensorflow-macos 0.1a0	tensorflow-macos 0.1a0	tensorflow 2.4

サンプルコードにはシンプルなmnistのクラス分類を使っています。

import tensorflow as tf
mnist = tf.keras.datasets.mnist
import time

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

start = time.perf_counter()
model.fit(x_train, y_train, epochs=5)
end = time.perf_counter()
print(f"Elapsed time : {end - start:.3f} s.")

model.evaluate(x_test, y_test)

実行結果

	Macbook Air (2020,M1)	Macbook Pro (2017,Intel)	開発に使っている学習用PC
5epochにかかった時間	6.268 s	23.876 s	8.053 s

Intel Macに対してM1 Macの方が4倍程度、速いという結果になりました。

素晴らしく速いですね。
たしかMacbook Pro(2017)の方がMacbook Air(2020)よりも値段が高かったような気がするんですが・・・。

さらにいつも開発に使っているRTX 2080よりも速かったのはもっと驚きの結果でした。

RTX 2080の方は特にチューニング作業はしていないので、使い方次第では実はもっと速くなる可能性はありますが、M1 Macだと何も気にせずこのぐらいの速度でトレーニングできるのは素晴らしいです。

これからDeep Learningをやりたくて個人で開発環境を整えるのであれば、下手に高いグラフィックボードを購入するよりは、M1が搭載されたMadbook Airを買った方がコスパは良いと思います。

追記

@nowokayさんから「RTX 2080の環境はGPUのコアを計算処理にフルに使っていないために、データ転送がボトルネックになっているのではないか」という指摘をいただきました。

実際に上記のスクリプトの実行中にRTX 2080の環境でnvidia-smiをして確認したところ、GPUの使用率が15%でした。
そこでアドバイスの通りにレイヤーをConv2Dに変更したり、フィルタ数を大きくすることでRTX 2080の環境でのGPUの使用率を95%まで上げ、再度比較をしました。

import tensorflow as tf
mnist = tf.keras.datasets.mnist
import time

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)
x_test  = x_test.reshape(x_test.shape[0], 28, 28, 1)

model = tf.keras.models.Sequential([
  tf.keras.layers.Conv2D(1024, (3,3), activation='relu', input_shape=(28, 28, 1)),
  tf.keras.layers.MaxPooling2D(2,2),
  tf.keras.layers.Conv2D(512, (3,3), activation='relu'),
  tf.keras.layers.MaxPooling2D(2,2),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(256, activation='relu'),
  tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

start = time.perf_counter()
model.fit(x_train, y_train, epochs=5)
end = time.perf_counter()
print("Elapsed time : {0} s.".format(end-start))

model.evaluate(x_test, y_test)

(※ Macbook Pro 2017はメルカリで売ってしまったので比較ができませんでした。)

	Macbook Air (2020,M1)	開発に使っている学習用PC
5epochにかかった時間	2629.628 s	135.903s

RTX 2080の方が圧倒的に短い時間で学習できました。

それぞれの環境の性能の違いがどれくらいなのかは、どこがボトルネックになっているかなど正確に比較する必要がありそうです。
実行時間の比較については参考程度にしてください。

@nowokayさんに指摘されるまでは不適切な条件での実験になっていました。
申し訳ありませんでした。そしてご指摘ありがとうございました！

考察、議論

GPU? CPU? Neural Engine?

tensorflow-macでは、
以下のように使用するデバイスを選択するオプションが用意されています。

from tensorflow.python.compiler.mlcompute import mlcompute
mlcompute.set_mlc_device(device_name="cpu")

# cpu : CPUだけを利用する。
# gpu : GPUだけを利用する。
# any : 最適なデバイスを使用する。(デフォルト)

Intel Macの場合、速い順にany > gpu > cpu でした。
cpuよりgpuの方が速いのは、Intel Iris Plus Graphicsを使っているからだと思いますが、
gpuよりanyの方が速いのは、「最適なデバイス」を使っているからなのでしょうか。

M1 Macの場合、速い順にany > cpu > gpuでした。
gpuが遅いのはEagar Mode (Define by run)に対応していないだけなので、Eagar modeをオフにすれば速くなるかもしれません。
anyが速いのは、やはり「最適なデバイス」を使っているからで、FowardにはNeural Engineを、BackwardにはGPUを使うなどの工夫がしてあるのかな？などと考察しています。

これらのデバイスが同じメモリを共有しているのがM1の大きな特徴の一つですし。

どこかに詳しい記事がないかを探してみます。

参考文献

• tensorflow-macosのGithubリポジトリ
• Accelerating TensorFlow Performance on Mac (Tensorflow Blog)
• macOS に最適化された TensorFlow を Apple M1 チップ 環境で import するまでの Tips Python
• Apple Silicon M1 は自然言語処理も、ちょっと速いよ