Mac M1 Pro上TensorFlow执行model.fit出现NotFoundError求助

问题重现 #

已成功安装TensorFlow内核并导入Keras相关包，运行以下简单线性回归代码：

import tensorflow as tf

# 创建简单线性回归模型
model = tf.keras.Sequential()
model.add(tf.keras.layers.Dense(1, input_shape=(1,)))

model.compile(optimizer='sgd', loss='mean_squared_error')

x_train = [1, 2, 3, 4]
y_train = [2, 4, 6, 8]

model.fit(x_train, y_train, epochs=10)

仅执行最后一行model.fit时触发以下错误：

Epoch 1/10
2023-01-18 15:08:48.858121: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:114] Plugin optimizer for device_type GPU is enabled.
2023-01-18 15:08:48.966479: 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: 0x165b6cbb0
2023-01-18 15:08:48.966508: 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: 0x165b6cbb0
---------------------------------------------------------------------------
NotFoundError                             Traceback (most recent call last)
Cell In[15], line 12
      9 x_train = [1, 2, 3, 4]
     10 y_train = [2, 4, 6, 8]
---> 12 model.fit(x_train, y_train, epochs=10)

File ~/miniconda3/envs/tensorflow/lib/python3.10/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # 要查看完整堆栈跟踪，请调用:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File ~/miniconda3/envs/tensorflow/lib/python3.10/site-packages/tensorflow/python/eager/execute.py:52, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     50 try:
     51   ctx.ensure_initialized()
---> 52   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     53                                       inputs, attrs, num_outputs)
     54 except core._NotOkStatusException as e:
     55   if name is not None:

NotFoundError: Graph execution error:
...
    File "/Users/Vetle/miniconda3/envs/tensorflow/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_1'
could not find registered platform with id: 0x165b6cbb0

当前环境信息 #

• Python平台：macOS-13.0-arm64-arm-64bit
• TensorFlow版本：2.11.0
• Keras版本：2.11.0
• Python版本：3.10.8 | conda-forge打包 | (main, Nov 22 2022, 08:25:29) [Clang 14.0.6 ]
• Pandas版本：1.5.2
• Scikit-Learn版本：1.2.0
• GPU可用

解决方案 #

1. 禁用XLA加速 #

在代码开头添加以下配置，强制TensorFlow关闭XLA编译：

import tensorflow as tf
tf.config.set_visible_devices([], 'GPU')
tf.config.optimizer.set_jit(False)

2. 调整TensorFlow版本 #

• 降级到2.10.0：TensorFlow 2.11在macOS ARM平台的XLA支持存在兼容性问题，降级到2.10版本可解决：

  conda install tensorflow=2.10.0
  

• 升级到2.12及以上：官方后续版本修复了ARM平台的部分兼容性问题，尝试升级到最新稳定版：

  conda install tensorflow>=2.12.0
  

3. 检查conda环境架构 #

确保conda环境为arm64架构，避免混用x86_64包：

conda info

查看输出中platform字段应为osx-arm64，若为x86_64则需重新创建适配arm64的环境。

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