概要

ビルド済みの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

Fashion MNIST

CIFAR-10

まとめ

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