TensorFlow自定义Transformer训练时遇OperatorNotAllowedInGraphError错误求助

看起来你遇到的这个OperatorNotAllowedInGraphError问题，核心原因是在TensorFlow的Graph模式（调用fit时默认启用）下，你在模型的call方法里做了一些Graph不支持的操作，另外还有几个数据和配置不匹配的小问题，我来一步步帮你梳理解决：

问题1：Graph模式下直接操作符号张量 #

当模型通过fit训练时，TensorFlow会自动切换到Graph模式，此时call方法里的inputs是符号tf.Tensor，不能像普通Python对象那样直接做这些操作：

• 用普通print()打印类型：这会触发对符号张量的迭代/内部访问，Graph模式不允许；
• 直接解包feature, label = inputs：虽然理论上符号元组可以解包，但如果操作不当也会触发错误，加上你后续没有用到label，反而容易出问题。

解决方法：

• 调试用tf.print()替代print()，它是Graph兼容的，可以安全打印张量信息；
• 解包操作如果需要保留，确保是在Graph允许的场景下，或者可以通过索引方式获取元素（比如feature = inputs[0]）。

问题2：标签与损失函数不兼容 #

你用了BinaryCrossentropy二分类损失，但你的标签是字符串'A'/'B'，这完全不匹配——损失函数需要接收数值型标签。

解决方法：
把字符串标签转换成0/1的数值，比如'A'→0，'B'→1，同时确保形状符合模型输出要求。

问题3：数据集变量名错误+重复设置batch_size #

你的代码里定义的数据集变量是dataset，但model.fit里写的是train_data，这会导致找不到数据集；另外dataset已经通过batch(1)做了批处理，不需要在fit里再重复指定batch_size=1。

修正后的完整代码示例 #

import numpy as np
import tensorflow as tf

# 1. 预处理标签：将字符串转为二分类数值
batched_features = tf.constant([[[1, 3], [2, 3]],
                                 [[2, 1], [1, 2]],
                                 [[3, 3], [3, 2]]], shape=(3, 2, 2))
# 把'A'转0，'B'转1，同时保持原形状(3,2,1)
batched_labels = tf.constant([['A', 'A'], ['B', 'B'], ['A', 'B']])
batched_labels = tf.where(batched_labels == 'A', 0, 1)
batched_labels = tf.expand_dims(batched_labels, axis=-1)

# 构建数据集
dataset = tf.data.Dataset.from_tensor_slices((batched_features, batched_labels))
dataset = dataset.batch(1)

# 可选：验证数据集格式
for element in dataset.as_numpy_iterator():
    print(element)

class MyTransformer(tf.keras.Model):
    def __init__(self):
        super().__init__()
        # 这里后续可以添加你的Transformer核心层，比如MultiHeadAttention、FeedForward等
        # 先添加一个适配二分类的输出层
        self.output_layer = tf.keras.layers.Dense(1, activation='sigmoid')

    def call(self, inputs, training=False):
        # Graph兼容的调试打印（可选）
        tf.print("Input tensor type:", type(inputs))
        # 安全解包输入张量
        feature, _ = inputs  # 这里label暂时用不到，用_占位
        # 通过输出层生成符合二分类要求的预测结果
        return self.output_layer(feature)

model = MyTransformer()
# 配置模型：确保损失函数和输出匹配（因为用了sigmoid，from_logits设为False）
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
              loss=tf.keras.losses.BinaryCrossentropy(from_logits=False),
              metrics=[tf.keras.metrics.BinaryAccuracy(),
                       tf.keras.metrics.FalseNegatives()])

# 启动训练：用正确的数据集变量，去掉重复的batch_size参数
model.fit(dataset, epochs=1)

关键修正说明 #

• 标签转换：将字符串标签转为数值，让损失函数可以正常计算；
• 模型输出层：添加了Dense(1, activation='sigmoid')，确保输出符合二分类任务的要求；
• Graph兼容操作：用tf.print替代普通print，避免触发符号张量的非法访问；
• 配置修正：修正了数据集变量名，移除重复的batch_size设置。

备注：内容来源于stack exchange，提问作者Mihail Vladov