M1 Pro设备TensorFlow GPU训练报错求助：找不到注册平台ID

我在M1/M2 Mac上安装了带GPU支持的TensorFlow，尝试将原有循环神经网络项目改为GPU运行，未修改核心代码，但训练时触发如下报错。

模型与训练代码 #

# Recurent NN
model = tf.keras.Sequential([
    tf.keras.layers.LSTM(64, return_sequences=True, input_shape=(len(data[0]), 1)),
    tf.keras.layers.Dropout(0.25),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.LSTM(32),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(output_dense, activation='softmax')
])

print(model.summary())


# Compile the model
model.compile(optimizer=optimizer, loss=loss_function, metrics=['accuracy'])

# Train the model
history = model.fit(data_train, labels_train, batch_size=batch_size, epochs=epochs, validation_data=(data_test, labels_test))

报错信息 #

2023-03-09 20:53:23.670721: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:418 : NOT_FOUND: could not find registered platform with id: 0x143695190
Traceback (most recent call last):
  File "/Users/kaanerbay/PycharmProjects/tftest/main.py", line 13, in <module>
    model.fit(x_train, y_train, epochs=5, batch_size=64)
  File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/tensorflow/python/eager/execute.py", line 52, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.NotFoundError: Graph execution error:

Detected at node 'StatefulPartitionedCall_212' defined at (most recent call last):
    File "/Users/kaanerbay/PycharmProjects/tftest/main.py", line 13, in <module>
      model.fit(x_train, y_train, epochs=5, batch_size=64)
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/engine/training.py", line 1650, in fit
      tmp_logs = self.train_function(iterator)
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/engine/training.py", line 1249, in train_function
      return step_function(self, iterator)
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/engine/training.py", line 1233, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/engine/training.py", line 1222, in run_step
      outputs = model.train_step(data)
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/engine/training.py", line 1027, in train_step
      self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/optimizers/optimizer_experimental/optimizer.py", line 527, in minimize
      self.apply_gradients(grads_and_vars)
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/optimizers/optimizer_experimental/optimizer.py", line 1140, in apply_gradients
      return super().apply_gradients(grads_and_vars, name=name)
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/optimizers/optimizer_experimental/optimizer.py", line 634, in apply_gradients
      iteration = self._internal_apply_gradients(grads_and_vars)
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/optimizers/optimizer_experimental/optimizer.py", line 1166, in _internal_apply_gradients
      return tf.__internal__.distribute.interim.maybe_merge_call(
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/optimizers/optimizer_experimental/optimizer.py", line 1216, in _distributed_apply_gradients_fn
      distribution.extended.update(
    File "/Users/kaanerbay/miniconda3/envs/main/lib/python3.10/site-packages/keras/optimizers/optimizer_experimental/optimizer.py", line 1211, in apply_grad_to_update_var
      return self._update_step_xla(grad, var, id(self._var_key(var)))
Node: 'StatefulPartitionedCall_212'
could not find registered platform with id: 0x143695190
     [[{{node StatefulPartitionedCall_212}}]] [Op:__inference_train_function_23355]

• 强制绑定Metal GPU后端：在代码开头添加以下配置，确保TensorFlow正确识别并使用Metal平台：

  import tensorflow as tf
  tf.keras.backend.clear_session()
  # 先清空原有设备配置，再绑定Metal GPU
  tf.config.set_visible_devices([], 'GPU')
  metal_devices = tf.config.list_physical_devices('Metal')
  if metal_devices:
      tf.config.set_visible_devices(metal_devices[0], 'GPU')
  

• 调整环境变量：设置环境变量避免Metal兼容性冲突，可在代码开头添加：

  import os
  os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
  

• 降级TensorFlow版本：部分新版本存在Metal平台注册bug，尝试切换到稳定兼容版本：

  pip uninstall tensorflow-macos tensorflow-metal
  pip install tensorflow-macos==2.11 tensorflow-metal==0.7.1
  

• 统一数据类型：Metal对数据类型要求严格，确保输入数据为float32：

  data_train = data_train.astype('float32')
  labels_train = labels_train.astype('float32')
  data_test = data_test.astype('float32')
  labels_test = labels_test.astype('float32')
  

• 临时禁用XLA编译：若为XLA与Metal冲突导致，可临时关闭XLA验证：

  os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices=false'
  

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