使用TensorFlow DirectML插件触发同优先级多OpKernel注册匹配报错

云上老周

2026-8-26

问题描述

硬件与需求：配备AMD RX Vega 64显卡，为缩短CPU训练模型的耗时，计划调用GPU完成模型训练，最终选定微软推出的DirectML方案。
前置操作：参照微软官方TensorFlow DirectML插件教程完成环境部署，当前使用Python 3.10.2版本，符合文档标注的版本支持要求；激活虚拟环境完成插件安装后，在Jupyter Notebook中运行数据集加载代码：

train_set = image_dataset_from_directory(
    train_path,
    labels="inferred",
    label_mode = 'categorical',
    image_size=IMG_SIZE,
    batch_size=batch_size,
)

dev_set = image_dataset_from_directory(
    dev_path,
    labels="inferred",
    label_mode = 'categorical',
    image_size=IMG_SIZE,
    batch_size=batch_size,
)

test_set = image_dataset_from_directory(
    test_path,
    labels="inferred",
    label_mode = 'categorical',
    image_size=IMG_SIZE,
    batch_size=batch_size,
)
class_names = train_set.class_names
num_classes = len(class_names)

报错信息：运行上述代码后触发如下错误：

InvalidArgumentError: Multiple OpKernel registrations match NodeDef at the same priority '{{node Equal}}': 'op: "Equal" device_type: "GPU" constraint { name: "T" allowed_values { list { type: DT_INT64 } } }' and 'op: "Equal" device_type: "GPU" constraint { name: "T" allowed_values { list { type: DT_INT64 } } }' [Op:Equal]

环境检测结果：执行如下环境校验代码

print(tf.__version__)
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

输出显示TensorFlow版本为2.9.1，检测到可用GPU数量为1，已正常识别RX Vega 64显卡，但无法定位报错原因，也不能确认当前环境是否可正常使用DirectML完成模型训练。

问题原因

该报错由TensorFlow GPU算子重复注册导致：当前环境同时存在两套GPU侧OpKernel（算子内核）注册信息——一套来自已安装的TensorFlow自带的原生GPU算子实现，另一套来自DirectML插件提供的算子实现，两套注册信息优先级完全一致，TensorFlow执行算子调度时无法选定唯一实现，直接抛出参数错误。
现有环境已经可以正常识别GPU，说明硬件驱动、DirectML基础组件没有问题，只需要解决包版本冲突即可正常使用。

解决步骤

完全清理当前虚拟环境内所有TensorFlow相关安装包，避免残留文件导致冲突，执行命令：
pip uninstall -y tensorflow tensorflow-gpu tensorflow-directml tensorflow-directml-plugin keras
安装匹配版本的依赖包，注意不要安装带原生GPU支持的tensorflow版本，选择CPU版本搭配DirectML插件即可：
pip install tensorflow-cpu==2.10.0 tensorflow-directml-plugin==0.2.0.dev221012
完全关闭Jupyter Notebook并重启内核，不要复用之前的运行时进程。
先运行简单GPU计算测试验证环境可用性，测试代码如下：

import tensorflow as tf
print(f"TensorFlow版本: {tf.__version__}")
print(f"识别到的GPU设备: {tf.config.list_physical_devices('GPU')}")
# 执行GPU侧计算测试
with tf.device('/GPU:0'):
    mat_a = tf.random.normal((1000, 1000))
    mat_b = tf.random.normal((1000, 1000))
    mat_c = tf.matmul(mat_a, mat_b)
print("GPU计算测试通过，环境可用")

测试无报错后，再运行之前的数据集加载代码即可正常执行。

注意：不要在该环境中安装CUDA、cuDNN相关的TensorFlow GPU依赖，这类依赖会再次引入重复的算子注册，触发同类报错。

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