Tensorflow 在 Mac Studio M1 Ultra 上非常慢

问：

我有一台 Mac Studio M1 Ultra，我正在尝试训练一个简单的 RNN 来预测时间序列。

代码如下

import numpy as np
import keras 
import tensorflow as tf

def generate_time_series(batch_size, n_steps):
    freq1, freq2, offsets1, offsets2 = np.random.rand(4, batch_size, 1)
    time = np.linspace(0, 1, n_steps)
    series = 0.5 * np.sin((time - offsets1) * (freq1 * 10 + 10)) # wave 1
    series += 0.2 * np.sin((time - offsets2) * (freq2 * 20 + 20)) # + wave 2
    series += 0.1 * (np.random.rand(batch_size, n_steps) - 0.5) # + noise
    return series[..., np.newaxis].astype(np.float32)

np.random.seed(42)
n_steps = 50
series = generate_time_series(10000, n_steps + 1)

X_train, y_train = series[:7000, :n_steps], series[:7000, -1]
X_valid, y_valid = series[7000:9000, :n_steps], series[7000:9000, -1]
X_test, y_test = series[9000:, :n_steps], series[9000:, -1]

# Implementing simple RRN
np.random.seed(42)
tf.random.set_seed(42)

model = keras.models.Sequential([keras.layers.SimpleRNN(1, input_shape=[None, 1])])

optimizer = tf.keras.optimizers.legacy.Adam(learning_rate=0.005)
model.compile(loss="mse", optimizer=optimizer)
history = model.fit(X_train, y_train, epochs=20,
                    validation_data=(X_valid, y_valid))

在我的 Mac M1 上，我按照此处的说明安装了 2.13.0，我有 2.13.1。Python 版本为 3.11。tensorflowkeras

当我运行代码时，我在训练期间得到以下输出

2023-11-14 17:55:24.172843: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:114] Plugin optimizer for device_type GPU is enabled.
Epoch 1/20
219/219 [==============================] - 148s 676ms/step - loss: 0.4310 - val_loss: 0.2155
Epoch 2/20
219/219 [==============================] - 143s 651ms/step - loss: 0.1627 - val_loss: 0.1514
Epoch 3/20
219/219 [==============================] - 142s 649ms/step - loss: 0.1462 - val_loss: 0.1488
Epoch 4/20
219/219 [==============================] - 141s 644ms/step - loss: 0.1474 - val_loss: 0.1475
Epoch 5/20
219/219 [==============================] - 142s 649ms/step - loss: 0.1477 - val_loss: 0.1508
Epoch 6/20
219/219 [==============================] - 142s 649ms/step - loss: 0.1006 - val_loss: 0.0617
...

这表明每个 epoch 需要 ~140 秒。

但是，如果我在 Google Colaboratory 上运行相同的代码，我会得到

Epoch 1/20
219/219 [==============================] - 6s 24ms/step - loss: 0.0485 - val_loss: 0.0173
Epoch 2/20
219/219 [==============================] - 3s 13ms/step - loss: 0.0132 - val_loss: 0.0118
Epoch 3/20
219/219 [==============================] - 4s 17ms/step - loss: 0.0119 - val_loss: 0.0113
Epoch 4/20
219/219 [==============================] - 4s 18ms/step - loss: 0.0116 - val_loss: 0.0110
Epoch 5/20
219/219 [==============================] - 5s 23ms/step - loss: 0.0114 - val_loss: 0.0109
Epoch 6/20
219/219 [==============================] - 2s 8ms/step - loss: 0.0114 - val_loss: 0.0109
Epoch 7/20
219/219 [==============================] - 2s 8ms/step - loss: 0.0114 - val_loss: 0.0109

速度快 100 倍！

为什么会有这样的差异？

我读过 tensorflow 在 Silicon Apple 上存在一些性能问题。有什么方法可以达到这样的性能吗？