TensorFlowをCPUとGPUで比較する

TensorFlowを安価でなるべく速く実行するにはどのような環境がよいのかを自分なりに検討してみました．

測定方法

以下のスクリプトを実行して，速度を比較します．環境変数CUDA_VISIBLE_DEVICESでGPUの利用するかを指定します．環境変数TF_ENABLE_ONEDNN_OPTSでoneDNNを使用するかを指定します．

requirements.txt

numpy==1.22.4
tensorflow==2.9.1
protobuf==3.19.4

mnist.py

keras.datasets.mnistを使用した簡単な機械学習のスクリプトです．

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "1"

import numpy as np
from tensorflow import keras, config
from tensorflow.keras import layers

gpus = config.list_physical_devices(device_type = 'GPU')
if len(gpus)>0:
    print(f">> GPU detected. {gpus[0].name}")
    config.experimental.set_memory_growth(gpus[0], True)


# https://keras.io/examples/vision/mnist_convnet/

# Model / data parameters
num_classes = 10
input_shape = (28, 28, 1)

# the data, split between train and test sets
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# Scale images to the [0, 1] range
x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255
# Make sure images have shape (28, 28, 1)
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)
print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")


# convert class vectors to binary class matrices
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

model = keras.Sequential(
    [
        keras.Input(shape=input_shape),
        layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Flatten(),
        layers.Dropout(0.5),
        layers.Dense(num_classes, activation="softmax"),
    ]
)

model.summary()

batch_size = 1024
epochs = 4

model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)

score = model.evaluate(x_test, y_test, verbose=0)
print("Test loss:", score[0])
print("Test accuracy:", score[1])

実行結果

所有しているパソコンや，Google Colabを使用して測定しました．

CPU Xeon E5-2673 v4, oneDNN ON (Colab)

lscpuコマンドで，CPU family: 6, Model: 79, Model name: Intel(R) Xeon(R) CPU @ 2.20GHzと表示されました．1コア2スレッド， 2.2GHzの割当です．

Epoch 1/4
53/53 [==============================] - 38s 712ms/step - loss: 1.0875 - accuracy: 0.6790 - val_loss: 0.2413 - val_accuracy: 0.9358
Epoch 2/4
53/53 [==============================] - 39s 732ms/step - loss: 0.2834 - accuracy: 0.9144 - val_loss: 0.1363 - val_accuracy: 0.9630
Epoch 3/4
53/53 [==============================] - 38s 714ms/step - loss: 0.1865 - accuracy: 0.9437 - val_loss: 0.1003 - val_accuracy: 0.9730
Epoch 4/4
53/53 [==============================] - 37s 702ms/step - loss: 0.1479 - accuracy: 0.9554 - val_loss: 0.0834 - val_accuracy: 0.9775
Test loss: 0.08833741396665573
Test accuracy: 0.9740999937057495

遅いです．

CPU i5-12400, oneDNN ON

CPU i5-12400は，6コア12スレッド，周波数は2.50 GHzです．

Epoch 1/4
53/53 [==============================] - 5s 88ms/step - loss: 1.1321 - accuracy: 0.6795 - val_loss: 0.2417 - val_accuracy: 0.9370
Epoch 2/4
53/53 [==============================] - 5s 89ms/step - loss: 0.2765 - accuracy: 0.9158 - val_loss: 0.1314 - val_accuracy: 0.9645
Epoch 3/4
53/53 [==============================] - 5s 86ms/step - loss: 0.1785 - accuracy: 0.9461 - val_loss: 0.0968 - val_accuracy: 0.9752
Epoch 4/4
53/53 [==============================] - 5s 86ms/step - loss: 0.1406 - accuracy: 0.9575 - val_loss: 0.0807 - val_accuracy: 0.9785
Test loss: 0.08474256843328476
Test accuracy: 0.9753999710083008

コア数が6倍になり，ターボブーストしているため，8.3倍高速になりました．

CPU i5-12400, oneDNN OFF

Epoch 1/4
53/53 [==============================] - 7s 127ms/step - loss: 1.1301 - accuracy: 0.6792 - val_loss: 0.2377 - val_accuracy: 0.9377
Epoch 2/4
53/53 [==============================] - 7s 123ms/step - loss: 0.2693 - accuracy: 0.9194 - val_loss: 0.1267 - val_accuracy: 0.9682
Epoch 3/4
53/53 [==============================] - 7s 123ms/step - loss: 0.1752 - accuracy: 0.9478 - val_loss: 0.0956 - val_accuracy: 0.9740
Epoch 4/4
53/53 [==============================] - 7s 123ms/step - loss: 0.1390 - accuracy: 0.9583 - val_loss: 0.0787 - val_accuracy: 0.9788
Test loss: 0.08312420547008514
Test accuracy: 0.9757000207901001

CPUを使用するならば，oneDNNにより1.4倍ほど高速になりました．

GPU GT1030

GeForce GT 1030のCUDAコアは384個で，周波数は1.4GHzです*1．

Epoch 1/4
53/53 [==============================] - 6s 82ms/step - loss: 1.1407 - accuracy: 0.6659 - val_loss: 0.2556 - val_accuracy: 0.9283
Epoch 2/4
53/53 [==============================] - 4s 71ms/step - loss: 0.2913 - accuracy: 0.9130 - val_loss: 0.1345 - val_accuracy: 0.9645
Epoch 3/4
53/53 [==============================] - 4s 71ms/step - loss: 0.1850 - accuracy: 0.9448 - val_loss: 0.0993 - val_accuracy: 0.9745
Epoch 4/4
53/53 [==============================] - 4s 71ms/step - loss: 0.1468 - accuracy: 0.9561 - val_loss: 0.0825 - val_accuracy: 0.9775
Test loss: 0.0853961631655693
Test accuracy: 0.974399983882904

nvidia-smiコマンドでGPU使用率は97%，メモリ使用量は1Gbyteでした．

i5-12400の2.5GHz * 6 core = 15, GT 1030の 1.4GHz * 384 core = 537.6 を比較すると36倍高速になりそうな気がしますが，CPUのほうが1クロックでたくさんの処理ができるので，結局はCPUと同じぐらいの速度になるようです．GPUならばとりあえず高速だというイメージだったので，CPUがいい勝負をしているのは意外でした．

GPU Tesla T4 (Colab)

Tesla T4 のCUDAコアは2560個です*2．

Epoch 1/4
53/53 [==============================] - 13s 29ms/step - loss: 1.1104 - accuracy: 0.6804 - val_loss: 0.2412 - val_accuracy: 0.9353
Epoch 2/4
53/53 [==============================] - 1s 20ms/step - loss: 0.2780 - accuracy: 0.9169 - val_loss: 0.1299 - val_accuracy: 0.9667
Epoch 3/4
53/53 [==============================] - 1s 21ms/step - loss: 0.1759 - accuracy: 0.9483 - val_loss: 0.0967 - val_accuracy: 0.9767
Epoch 4/4
53/53 [==============================] - 1s 20ms/step - loss: 0.1415 - accuracy: 0.9579 - val_loss: 0.0813 - val_accuracy: 0.9780
Test loss: 0.08577422052621841
Test accuracy: 0.9743000268936157

nvidia-smiコマンドでGPU使用率は76%でした．

CUDAコアが6.6倍になると，3.55倍高速になりました．GPU使用率が下がっているので，GPUの能力を十分に使用できていない可能性があります．

考察

最新のCPUならば，oneDNNを有効にするだけで高速になります．i5-12400は3万円，GT1030は1万円ぐらいで購入でき，同じような速度でした．Tesla T4と同じCUDAコアの数を搭載したRTX 3050は5万円ぐらいで購入できます．

しかし，GPUにはデメリットが少なくとも3個あります．まず，ドライバーのインストールなど環境を整えるのが面倒です．次に，GPUで高速に実行できるようにソフトウェアをある程度調整しなければならないことです．最後に，メモリーを自由に増設できないことです．

mnistでもGPUメモリーを1Gbyteほど使用するので，複雑なモデルの機械学習を実行するならば，GPUメモリーが不足する可能性があります．GPUメモリーへの転送はとても時間がかかる処理なので，GPUメモリーが不足するとせっかくのGPUのパワーを活かせないのです．GeForce最強GPUであるRTX 3090 でもGPUメモリーは24Gbyteです．Tesla A100ならば80GBのGPUメモリーを扱えます． CPUならば安価でよりたくさんのメモリを扱えます．i5-12400ならば最大で128 GBです．

ひとまずCPUで実行してみて，とても時間がかかるならば，必要なスペック(特にCUDAコア数とGPUメモリー)のGPUをクラウドで借りるのが現実的な気がしました．