Keras 3.0になってマルチバックエンド（TensorFlow,jax, PyTorch)が使えるようになったので速度を比較してみる(2) MNIST（手書き数字認識）でPythonでCNN（畳み込みニューラルネットワーク）でやってみる。TensorFlowが一番速く、PyTorch, JAXの順（CPUのみ使用の場合）。

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

https://keras.io/keras_3/

前回はDNN（Deep Neural Networks)で比較してみた。

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

今回はCNN（畳み込みニューラルネットワーク）でやってみよう。

初期設定は共通で、Kerasのサイトの例題をFunctional APIに書き直して使っている。

import os

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

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

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

import keras

import numpy as np

#　データは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

x_train = np.expand_dims(x_train, -1)

x_test = np.expand_dims(x_test, -1)

# Functional APIでCNN設定

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

x = keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu")(inputs)

x = keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu")(x)

x = keras.layers.MaxPooling2D(pool_size=(2, 2))(x)

x = keras.layers.Conv2D(128, kernel_size=(3, 3), activation="relu")(x)

x = keras.layers.Conv2D(128, kernel_size=(3, 3), activation="relu")(x)

x = keras.layers.GlobalAveragePooling2D()(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"),])

時間を計測する学習部分はこのコード。

batch_size = 128

epochs = 5

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がやっぱり速いのは前と同じだが、今度はJAXが非常に遅くなった。

謎だ…GPU使えば結果は変わるのかな。次はRNNをやってみよう。