Apple Silicon M1でのTensorFlowの動作速度検証

#はじめに

発売当初の絶賛ムードに対して、若干時間がたった現在では「そうでもないのでは？」的な記事も増えている気がするM1 Mac。
TensorFlowのM1対応に関しても、なんとも評価しづらいような結果の記事が出てきているが、ここでは筆者が購入したMac Miniでの簡単な検証結果を記事にする。

#基本的にGPUを使うと遅い

M1に最適化されたというTensorFlowでは、使用するデバイスを下記のように指定できる。

from tensorflow.python.compiler.mlcompute import mlcompute
mlcompute.set_mlc_device(device_name="gpu") # 'any' or 'cpu' or 'gpu'

デフォルトでは'any'で最適なデバイスを勝手に選んでくれるという。

VGG19をCIFAR10で学習させた場合の１エポックの学習時間は以下の通りだった。(バッチサイズは100)

any	cpu	gpu
247s	243s	458s

'gpu'が一番遅い。'any'と'cpu'が同じくらい。ResNet50でも似たような結果になる。
アクティビティモニタでGPUやCPUの履歴を見ると、GPUを使うのは'gpu'とした場合のみで、'any'としてもCPUしか使わないようだ。
ちなみに、CPUを使用している際にもCPUの使用率としては割と余裕があり100%使い切っていない(使いきれていないというべきか？)。

M1はNeural Engineというのも搭載しているが、私の理解ではこれは推論でしか使われないもので、今回のように訓練では関係ないと思われる。
つまり、'any'とした際にはCPUが'最適なデバイス'として選択されていることになるだろう。
実際GPUを使うと遅くなるわけだから、最適と言われればその通りだけれど、これでいいのか？という疑問は感じざるを得ない。今後のアップデートで改善されるかもしれないが。

#GPUを使った方が早い場合もある

以下のようなコードでモデルを作成すると、GPUとCPUの処理時間が逆転した。

def make_model():
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Conv2D(64, (3,3), padding='same', activation='relu',input_shape=(32,32,3)))
    model.add(tf.keras.layers.BatchNormalization())
    model.add(tf.keras.layers.MaxPooling2D((2,2), padding='same'))

    model.add(tf.keras.layers.Conv2D(128, (3,3), padding='same', activation='relu'))
    model.add(tf.keras.layers.BatchNormalization())
    model.add(tf.keras.layers.MaxPooling2D((2,2), padding='same'))

    model.add(tf.keras.layers.Conv2D(256, (3,3), padding='same', activation='relu'))
    model.add(tf.keras.layers.BatchNormalization())
    model.add(tf.keras.layers.MaxPooling2D((2,2), padding='same'))

    model.add(tf.keras.layers.Conv2D(512, (3,3), padding='same', activation='relu'))
    model.add(tf.keras.layers.BatchNormalization())
    model.add(tf.keras.layers.GlobalMaxPooling2D())
    model.add(tf.keras.layers.Dense(10, activation='softmax'))
    return model

any	cpu	gpu
43s	43s	35s

これはGPUが早い。Conv2Dの数を減らすとGPUが有利になるようだ。
GPUの方が早いので、'any'で最適なデバイスが選択できていないとも言える。わざわざモデルの内容をみてデバイスを決める処理が入っているとも思えないので、現状常に'any'ではCPUが選択されるのだろう。

#Intel CPUに比べれば速い(はず)

筆者はIntel版Macを持っていないので、WindowsPCとGoogleColabのGPUなしでの比較とする。
上にあげた、簡単なモデルで訓練を実行した結果がこちら。

デバイス	時間(s)
M1 Mac (CPU)	43
M1 Mac (GPU)	35
Google Colab (Xeon 2.20GHz 2cores)	537
Windows(i7-4710MQ 2.50GHz 4cores)	505

CPUでの処理速度としては圧倒的に速かった。
Intel版Macと比較しても、それなりに速いのではなかろうかと推測する。

#NvidiaのGPUと比べると厳しい

上記と同じ実験をGoogleColabのP100で実施すると１エポックが5秒程度で終わる(初回は12秒ぐらいだが)。
K80ならもう少し時間がかかりそうだが、それでもM1よりは速いだろう。

Geforce GT730Mというモバイル向けのチップでCUDA対応のTensorFlowで実行すると104秒だった。
M1の方が速いのは立派だが、これは古いチップでcoreが384しかないので、最近のものならばM1より早くなるだろう。

#テストコード全文

import tensorflow as tf
import sys

def make_model():
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Conv2D(64, (3,3), padding='same', activation='relu',input_shape=(32,32,3)))
    model.add(tf.keras.layers.BatchNormalization())
    model.add(tf.keras.layers.MaxPooling2D((2,2), padding='same'))

    model.add(tf.keras.layers.Conv2D(128, (3,3), padding='same', activation='relu'))
    model.add(tf.keras.layers.BatchNormalization())
    model.add(tf.keras.layers.MaxPooling2D((2,2), padding='same'))

    model.add(tf.keras.layers.Conv2D(256, (3,3), padding='same', activation='relu'))
    model.add(tf.keras.layers.BatchNormalization())
    model.add(tf.keras.layers.MaxPooling2D((2,2), padding='same'))

    model.add(tf.keras.layers.Conv2D(512, (3,3), padding='same', activation='relu'))
    model.add(tf.keras.layers.BatchNormalization())
    model.add(tf.keras.layers.GlobalMaxPooling2D())
    model.add(tf.keras.layers.Dense(10, activation='softmax'))
    return model


def test(device_name='any', model_name='VGG19'):
    from tensorflow.python.compiler.mlcompute import mlcompute
    mlcompute.set_mlc_device(device_name=device_name)

    train_data, validation_data = tf.keras.datasets.cifar10.load_data()

    batch_size = 100

    ds_train = tf.data.Dataset.from_tensor_slices(train_data).repeat(50000)
    ds_train = ds_train.batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE)
    ds_validation = tf.data.Dataset.from_tensor_slices(validation_data)
    ds_validation = ds_validation.batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE)

    if model_name=='VGG19':
        from tensorflow.keras.applications import VGG19
        model = VGG19(weights=None, input_shape=(32,32,3), classes=10)
    elif model_name=='ResNet50V2':
        from tensorflow.keras.applications import ResNet50V2
        model = ResNet50V2(weights=None, input_shape=(32,32,3), classes=10)
    else:
        model = make_model()
        model.summary()

    steps_per_epoch = 50000//batch_size

    loss = tf.keras.losses.SparseCategoricalCrossentropy()
    acc = tf.keras.metrics.SparseCategoricalAccuracy(name='acc')
    model.compile(optimizer=tf.keras.optimizers.Adam(),
                loss=loss, metrics=[acc])
    model.fit(ds_train, epochs=3, steps_per_epoch=steps_per_epoch,
                validation_data=ds_validation)

test( sys.argv[1], sys.argv[2])

#まとめ

現状ではM1に最適化したTensorFlowはCPU処理の最適化がメインで、GPUは有効に使えていないようである。
それでもCPU処理のみを考えると、かなり速いのではないか。
ただし、TensorFlow関係では色々と不具合の報告もあるようなので、現状では積極的におすすめはしない。

#参考

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

機械学習の開発環境としてのMacBook Air（M1）