NVIDIA RTX5060 Laptop GPUを入手したのでKeras3を使ってバックエンドにPyTorchとJax（WSL2利用）を使ってMNISTをCNNで学習するときの時間を比較した。ついでにGoogle ColabのT4 GPUでも同じことをした。Jax、めちゃくちゃ早いし、5060はさすがに性能アップしている。

この前、ゲーミングノートPCを買った。

 家で使うPCをゲーミングノートPC、ASUS TUF Gaming A16に買い替えた。CPUはAMD Ryzen 9 8940HX、メモリ32GB、GPUはNVIDIA GeForce RTX 5060 Laptop GPU、SSD 1TB。ゲームをしたいわけでなくてNVIDIAの最新GPUで機械学習・数値計算やろうかと。

で早速GPUのパワーを試すことにした。例題は一番簡単なMNISTでCNNを使う。この辺を参考に。

Keras3.0を使えばバックエンドはPyTorchでもJaxでもTensorflowでも基本、同じコードにできる。ただ、Tensorflowは最新のGPUに対応してないっぽいので外す。JaxもWindows版はGPU使えないので全部WSL2にした。

CUDA13にしないとRTX5060動かないとかいろいろ罠があって何度もインストールやり直したりしてようやく動いた。

コードはPytorchをバックエンドにするならこんな感じで。

import os

# 【重要】kerasをインポートする「前」にバックエンドをPyTorchに設定します

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

import keras

import torch

import numpy as np

import time

# --- 1. 設定の確認 ---

print(f"Keras Version: {keras.__version__}")

print(f"Current Backend: {keras.config.backend()}")

# GPUがPyTorchから見えているか確認

if torch.cuda.is_available():

    print(f"GPU is available: {torch.cuda.get_device_name(0)}")

    device = "cuda"

else:

    print("GPU is NOT available. Using CPU.")

    device = "cpu"

# --- 2. データの準備 (MNIST) ---

# Keras 3はNumPyデータをそのまま扱えます

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

# データの正規化と形状変更

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

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

x_train = np.expand_dims(x_train, -1)

x_test = np.expand_dims(x_test, -1)

# --- 3. モデルの定義 ---

# PyTorchバックエンドでもKerasの構文(Sequential/Functional API)が使えます

model = keras.Sequential([

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

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

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

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

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

    keras.layers.Flatten(),

    keras.layers.Dropout(0.5),

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

])

# モデルのコンパイル

# optimizerやlossはKerasの標準APIを使用します

model.compile(

    optimizer="adam",

    loss="sparse_categorical_crossentropy",

    metrics=["accuracy"],

)

model.summary()

start_time = time.perf_counter()

# --- 4. 学習の実行 ---

# Keras 3は、デフォルトで利用可能な最適なデバイス（GPUがあればGPU）を自動的に使用します

print("\nStarting training...")

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

# --- 5. 評価 ---

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

end_time = time.perf_counter()

print("\nKeras with PyTorch backend - GPU : NVIDIA RTX 5060 Laptop GPU")

print(f"Total training and evaluation time: {end_time - start_time:.2f} seconds")

print(f"Test loss: {score[0]}")

print(f"Test accuracy: {score[1]}")

jaxを使うときは

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

にする。

では

Pytorch + RTX5060

Jax ＋ RTX5060

うわ、Jaxのほうが倍以上速い。

せっかくなんでGoogle ColabでT4 GPU (Tesla T4)でもやってみよう。

Pytorch + T4

Jax ＋ T4

こちらもJaxのほうが速い。またRTX5060の性能がいいこともわかった。CPUも含めた結果はこちら。

                  CNNによるMNIST計算時間

GPU　＼ Keras Backend	PyTorch	Jax
RTX 5060 Laptop GPU	51.07s	20.17s
T4 GPU (Google Colab)	72.37s	28.78s
CPU (Ryzen 9)	127.14s	254.68s