M1 Max上TensorFlow GPU训练远慢于CPU的问题求助

我运行了一个改编自机器学习教程的时序模型简化示例，发现一个异常情况：CPU训练速度比GPU快50倍以上。

示例代码 #

from pathlib import Path
import pandas as pd
import tensorflow as tf

tf.keras.utils.get_file(
    "ridership.tgz",
    "https://github.com/ageron/data/raw/main/ridership.tgz",
    cache_dir=".",
    extract=True
)
path = Path("datasets/ridership/CTA_-_Ridership_-_Daily_Boarding_Totals.csv")
df = pd.read_csv(path, parse_dates=["service_date"])
df = df.sort_values("service_date").set_index("service_date")
df = df.drop_duplicates()

rail_train = df["rail_boardings"]["2016-01":"2018-12"] / 1e6
rail_valid = df["rail_boardings"]["2019-01":"2019-05"] / 1e6

seq_length = 56

train_ds = tf.keras.utils.timeseries_dataset_from_array(
    rail_train.to_numpy(),
    targets=rail_train[seq_length:],
    sequence_length=seq_length,
    batch_size=32,
    shuffle=True,
    seed=42
)
valid_ds = tf.keras.utils.timeseries_dataset_from_array(
    rail_valid.to_numpy(),
    targets=rail_valid[seq_length:],
    sequence_length=seq_length,
    batch_size=32
)

tf.random.set_seed(42)  # 确保结果可复现
deep_model = tf.keras.Sequential([
    tf.keras.layers.SimpleRNN(32, return_sequences=True, input_shape=[None, 1]),
    tf.keras.layers.SimpleRNN(32, return_sequences=True),
    tf.keras.layers.SimpleRNN(32),
    tf.keras.layers.Dense(1)
])

GPU运行情况 #

执行代码：

with tf.device('/gpu:0'):
    opt = tf.keras.optimizers.legacy.SGD(learning_rate=0.01, momentum=0.9)
    deep_model.compile(loss=tf.keras.losses.Huber(), optimizer=opt, metrics=["mae"])
    deep_model.fit(train_ds, validation_data=valid_ds, epochs=10)

运行结果：

Epoch 1/10
2023-01-25 15:08:13.733000: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:114] Plugin optimizer for device_type GPU is enabled.
33/33 [==============================] - ETA: 0s - loss: 0.0306 - mae: 0.1614
2023-01-25 15:12:10.354167: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:114] Plugin optimizer for device_type GPU is enabled.
33/33 [==============================] - 241s 7s/step - loss: 0.0306 - mae: 0.1614 - val_loss: 0.0045 - val_mae: 0.0783
Epoch 2/10
33/33 [==============================] - 265s 8s/step - loss: 0.0082 - mae: 0.0985 - val_loss: 0.0118 - val_mae: 0.1358
Epoch 3/10
33/33 [==============================] - 243s 7s/step - loss: 0.0066 - mae: 0.0838 - val_loss: 0.0030 - val_mae: 0.0567
Epoch 4/10
33/33 [==============================] - 236s 7s/step - loss: 0.0046 - mae: 0.0631 - val_loss: 0.0022 - val_mae: 0.0455
Epoch 5/10
17/33 [==============>...............] - ETA: 1:52 - loss: 0.0045 - mae: 0.0609

CPU运行情况 #

执行代码：

with tf.device('/cpu:0'):
    opt = tf.keras.optimizers.legacy.SGD(learning_rate=0.01, momentum=0.9)
    deep_model.compile(loss=tf.keras.losses.Huber(), optimizer=opt, metrics=["mae"])
    deep_model.fit(train_ds, validation_data=valid_ds, epochs=10)

运行结果：

Epoch 1/10
2023-01-25 15:35:43.883427: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:114] Plugin optimizer for device_type GPU is enabled.
33/33 [==============================] - ETA: 0s - loss: 0.0046 - mae: 0.0654
2023-01-25 15:35:48.833485: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:114] Plugin optimizer for device_type GPU is enabled.
33/33 [==============================] - 7s 160ms/step - loss: 0.0046 - mae: 0.0654 - val_loss: 0.0020 - val_mae: 0.0417
Epoch 2/10
33/33 [==============================] - 4s 128ms/step - loss: 0.0049 - mae: 0.0705 - val_loss: 0.0019 - val_mae: 0.0394
Epoch 3/10
33/33 [==============================] - 4s 135ms/step - loss: 0.0043 - mae: 0.0633 - val_loss: 0.0020 - val_mae: 0.0386
Epoch 4/10
33/33 [==============================] - 4s 131ms/step - loss: 0.0036 - mae: 0.0521 - val_loss: 0.0019 - val_mae: 0.0353
Epoch 5/10
33/33 [==============================] - 4s 127ms/step - loss: 0.0038 - mae: 0.0550 - val_loss: 0.0021 - val_mae: 0.0407
Epoch 6/10
33/33 [==============================] - 4s 123ms/step - loss: 0.0037 - mae: 0.0550 - val_loss: 0.0018 - val_mae: 0.0328
Epoch 7/10
33/33 [==============================] - 4s 123ms/step - loss: 0.0041 - mae: 0.0618 - val_loss: 0.0026 - val_mae: 0.0532
Epoch 8/10
33/33 [==============================] - 5s 141ms/step - loss: 0.0037 - mae: 0.0559 - val_loss: 0.0018 - val_mae: 0.0328
Epoch 9/10
33/33 [==============================] - 4s 129ms/step - loss: 0.0037 - mae: 0.0549 - val_loss: 0.0031 - val_mae: 0.0570
Epoch 10/10
33/33 [==============================] - 4s 129ms/step - loss: 0.0036 - mae: 0.0553 - val_loss: 0.0023 - val_mae: 0.0432

补充尝试 #

我理解简单模型有时在CPU上表现更快，因此尝试堆叠15层tf.keras.layers.SimpleRNN(32, return_sequences=True)，但GPU与CPU的速度差距反而进一步拉大。

环境依赖版本 #

tensorflow-datasets       4.7.0                    pypi_0    pypi
tensorflow-deps           2.9.0                         0    apple
tensorflow-estimator      2.11.0                   pypi_0    pypi
tensorflow-macos          2.11.0                   pypi_0    pypi
tensorflow-metadata       1.11.0                   pypi_0    pypi
tensorflow-metal          0.7.0                    pypi_0    pypi

恳请帮忙解决M1 Max上TensorFlow GPU训练远慢于CPU的异常问题。

内容的提问来源于stack exchange，提问作者Amiel