Keras 3.0になってマルチバックエンド（TensorFlow,jax, PyTorch)が使えるようになったので速度を比較してみる(1) MNIST（手書き数字認識）でPythonコード書いてDNNでやってみる。TensorFlowが一番速く、jax, PyTorchの順（CPUのみ使用の場合）。ちょっと意外。

Kerasが3.0になってまた独立してマルチバックエンドTensorFlow, jax, PyTorch)が使えるようになっていた。

https://keras.io/keras_3/

同じPythonコードで自由に選べるということで速度比較をしてみよう。jaxに期待している。

まずは設定のコード（Functional API利用）はほぼ共通で、KERAS_BACKENDで選べるようになっている。

import os

os.environ["KERAS_BACKEND"] = "tensorflow"

# os.environ["KERAS_BACKEND"] = "jax"

# os.environ["KERAS_BACKEND"] = "torch"

import keras

#　データはMNIST

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

x_train = x_train.astype("float32") / 255

x_test = x_test.astype("float32") / 255

# Functional APIでDense層＋ドロップアウト層を3層にしたDNNを設定

hidden_dim = 500

inputs = keras.Input(shape=(28, 28))

x = keras.layers.Flatten()(inputs)

x = keras.layers.Dense(hidden_dim, activation="relu")(x)

x = keras.layers.Dropout(0.5)(x)

x = keras.layers.Dense(hidden_dim, activation="relu")(x)

x = keras.layers.Dropout(0.5)(x)

x = keras.layers.Dense(hidden_dim, activation="relu")(x)

x = keras.layers.Dropout(0.5)(x)

outputs = keras.layers.Dense(10, activation="softmax")(x)

#　モデルの設定

model = keras.Model(inputs=inputs, outputs=outputs)

model.compile(optimizer=keras.optimizers.Adam(), loss=keras.losses.SparseCategoricalCrossentropy(),    

            metrics=[keras.metrics.SparseCategoricalAccuracy(name="acc"),])

で速度を比較したのはfitの部分。

batch_size = 128

epochs = 20

keras.utils.set_random_seed(1)

model.fit(

    x_train,

    y_train,

    batch_size=batch_size,

    epochs=epochs,

    validation_split=0.15,

)

score = model.evaluate(x_test, y_test, verbose=0)

print(score)

ではやってみよう。

TensorFlow

jax

PyTorch

まとめるとこうなった。TensorFlowが一番速かった。ちょっと意外。モデルが単純すぎるのかな？あるいはCPUだけならこうなるとか。

もうちょっと違う例でやってみよう。