TensorFlowのAVX2, FMAの有無で性能の比較をする

Developers's blog

2019.05.16中村 昴
TensorFlowのAVX2, FMAの有無で性能の比較をする

概要

ビルド済みのTensorFlowをpipでインストールすると、CPUの命令セットのAVX2やFMAが有効にならない。 AVX2やFMAを用いると、演算速度や精度の向上が期待できる。

参考 MMX, SSE, AVX, FMA命令とは何かを明らかにしたい会

そこで今回は、ソースコードからビルドしてAVX2とFMAを有効にしたTensorFlowを用意し、ビルド済みのパッケージをインストールしたものと、速度や精度の比較を行った。

環境

  • プロセッサ Intel Core i7-5557U 3.1 GHz
  • メモリ 16GB 1867 MHz DDR3
  • macOS Sierra 10.12.6
  • Python 3.6.8
  • TensorFlow 1.13.1

方法

学習は以下のようなCNNで行った。 データセットはMNIST、Fashion MNIST、CIFAR-10を用いた。

import tensorflow as tf

# データセットの選択
dataset = tf.keras.datasets.mnist
# dataset = tf.keras.datasets.fashion_mnist
# dataset = tf.keras.datasets.cifar10

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

# shapeをモデルの入力に合わせる
if len(x_train.shape) == 3:
    x_train = x_train.reshape(x_train.shape + (1,));
    x_test = x_test.reshape(x_test.shape + (1,));

# CNNモデルの作成
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(32,
        kernel_size=(3, 3),
        activation=tf.nn.relu,
        input_shape=x_train.shape[1:]),
    tf.keras.layers.Conv2D(64, (3, 3), activation=tf.nn.relu),
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    tf.keras.layers.Dropout(0.25),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation=tf.nn.relu),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])
model.compile(optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy'])

# 訓練
model.fit(x_train, y_train, epochs=12)
# 評価
model.evaluate(x_test, y_test)

計測は以下のコマンドで行った。

$ time python cnn.py

結果

MNIST

AVX2, FMA 有 AVX2, FMA 無
学習時間 26m37s 28m32s 6.7%
精度 0.9939 0.9928 +0.0011

Fashion MNIST

AVX2, FMA 有 AVX2, FMA 無
学習時間 25m30s 27m59s 8.9%
精度 0.9218 0.9241 -0.0023

CIFAR-10

AVX2, FMA 有 AVX2, FMA 無
学習時間 32m00s 37m04s 13.7%
精度 0.7049 0.7034 +0.0015

まとめ

AVX2, FMAが有効であると実行時間が短くなった。一方、精度はあまり変わらなかった。

お問い合わせはこちらから