如何解决Keras/TensorFlow中的OperatorNotAllowedInGraphError错误？

问题背景 #

使用Keras 3/TensorFlow 2.16.1将官方2D半监督SimCLR示例适配3D图像分类任务时，触发错误：

OperatorNotAllowedInGraphError: Iterating over a symbolic tf.Tensor is not allowed

错误原因分析 #

• 缺失keras.ops模块导入：代码中大量使用ops.xxx符号张量操作，但未导入对应模块，导致张量处理逻辑异常
• 2D增强层不适配3D数据：RandomFlip、RandomTranslation等默认是2D层，直接作用于3D张量会破坏张量结构，引发符号张量遍历错误
• 数据集结构与代码解构不匹配：train_step中错误地将输入数据解构为(unlabeled_images, _), (labeled_images, labels)，但实际传入的数据集结构是单组(图像,标签)
• 编码器输出维度不匹配：编码器原输出为1维sigmoid结果，与投影头期望的128维特征输入冲突，导致张量形状错误
• 损失函数与任务不匹配：二分类任务错误使用多分类的SparseCategoricalCrossentropy，且线性探针输出维度与类别数不符

分步修复方案 #

1. 导入缺失的keras.ops模块 #

在代码开头添加：

from tensorflow.keras import ops  # 新增导入，支持符号张量操作

2. 替换2D增强层为3D版本 #

修改get_augmenter函数，使用3D专属数据增强层，并更新自定义翻转层：

def get_augmenter(min_area):
    zoom_factor = 1.0 - math.sqrt(min_area)
    return tf.keras.Sequential(
        [
            layers.Rescaling(1 / 255),
            layers.RandomFlip("horizontal", data_format="channels_last"),
            layers.RandomTranslation(zoom_factor / 2, zoom_factor / 2, zoom_factor / 2),  # 3D三方向平移
            layers.RandomZoom((-zoom_factor, 0.0), (-zoom_factor, 0.0), (-zoom_factor, 0.0)),  # 3D三方向缩放
            RandomGeometricTransform3D(),
        ]
    )

class RandomGeometricTransform3D(layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.seed_generator = tf.random.experimental.Generator.from_seed(1337)

    def get_config(self):
        config = super().get_config()
        return config

    def call(self, images, training=True):
        if training:
            images = tf.image.random_flip_left_right(images, seed=self.seed_generator.make_seeds(2)[0][0])
            images = tf.image.random_flip_up_down(images, seed=self.seed_generator.make_seeds(2)[0][1])
        return images

    def compute_output_shape(self, input_shape):
        return input_shape

3. 修正train_step的数据解构逻辑 #

调整train_step的参数解构，匹配实际传入的数据集结构，并修正fit调用传入正确的数据集：

def train_step(self, data):
    # 修正数据解构：传入的labeled_train_dataset结构为(图像, 标签)
    labeled_images, labels = data
    # 若需真实无标签数据，需传入train_dataset而非labeled_train_dataset，此处复用有标签数据做示例
    unlabeled_images = labeled_images
    
    images = ops.concatenate((unlabeled_images, labeled_images), axis=0)
    # 后续增强、训练逻辑保持不变
    augmented_images_1 = self.contrastive_augmenter(images, training=True)
    augmented_images_2 = self.contrastive_augmenter(images, training=True)
    with tf.GradientTape() as tape:
        features_1 = self.encoder(augmented_images_1, training=True)
        features_2 = self.encoder(augmented_images_2, training=True)
        projections_1 = self.projection_head(features_1, training=True)
        projections_2 = self.projection_head(features_2, training=True)
        contrastive_loss = self.contrastive_loss(projections_1, projections_2)
    gradients = tape.gradient(
        contrastive_loss,
        self.encoder.trainable_weights + self.projection_head.trainable_weights,
    )
    self.contrastive_optimizer.apply_gradients(
        zip(gradients, self.encoder.trainable_weights + self.projection_head.trainable_weights)
    )
    self.contrastive_loss_tracker.update_state(contrastive_loss)

    preprocessed_images = self.classification_augmenter(labeled_images, training=True)
    with tf.GradientTape() as tape:
        features = self.encoder(preprocessed_images, training=False)
        class_logits = self.linear_probe(features, training=True)
        probe_loss = self.probe_loss(labels, class_logits)
    gradients = tape.gradient(probe_loss, self.linear_probe.trainable_weights)
    self.probe_optimizer.apply_gradients(zip(gradients, self.linear_probe.trainable_weights))
    self.probe_loss_tracker.update_state(probe_loss)
    self.probe_accuracy.update_state(labels, class_logits)
    
    # 返回指标结果（必须添加，否则Keras无法追踪指标）
    return {m.name: m.result() for m in self.metrics}

修正fit调用：

pretraining_history = pretraining_model.fit(
    train_dataset, epochs=num_epochs, validation_data=test_dataset
)

4. 调整编码器输出维度 #

修改get_encoder，移除分类层，输出与投影头匹配的128维特征：

def get_encoder():
    return tensorflow.keras.Sequential(
        [
            layers.Conv3D(16, kernel_size=(3, 3, 3), activation="relu", input_shape=(10,10,10,1)),
            layers.MaxPooling3D(pool_size=(2, 2, 2)),
            layers.Conv3D(32, kernel_size=(3, 3, 3), activation="relu"),
            layers.MaxPooling3D(pool_size=(2, 2, 2)),
            layers.Dropout(0.5),
            layers.Flatten(),
            layers.Dense(128, activation="relu"),  # 输出128维特征，匹配投影头输入
            layers.Dropout(0.5),
        ],
        name="encoder",
    )

5. 匹配二分类任务的损失函数与指标 #

修改ContrastiveModel的初始化和编译逻辑：

class ContrastiveModel(tensorflow.keras.Model):
    def __init__(self):
        super().__init__()
        self.temperature = temperature
        self.contrastive_augmenter = get_augmenter(**contrastive_augmentation)
        self.classification_augmenter = get_augmenter(**classification_augmentation)
        self.encoder = get_encoder()
        self.projection_head = tensorflow.keras.Sequential(
            [
                tensorflow.keras.Input(shape=(width,)),
                layers.Dense(width, activation="relu"),
                layers.Dense(width),
            ],
            name="projection_head",
        )
        # 线性探针改为二分类输出
        self.linear_probe = tensorflow.keras.Sequential(
            [layers.Input(shape=(width,)), layers.Dense(1, activation="sigmoid")],
            name="linear_probe",
        )

    def compile(self, contrastive_optimizer, probe_optimizer, **kwargs):
        super().compile(**kwargs)
        self.contrastive_optimizer = contrastive_optimizer
        self.probe_optimizer = probe_optimizer
        # 二分类损失函数
        self.probe_loss = tensorflow.keras.losses.BinaryCrossentropy(from_logits=False)
        self.contrastive_loss_tracker = tensorflow.keras.metrics.Mean(name="c_loss")
        self.contrastive_accuracy = tensorflow.keras.metrics.SparseCategoricalAccuracy(name="c_acc")
        self.probe_loss_tracker = tensorflow.keras.metrics.Mean(name="p_loss")
        # 二分类准确率指标
        self.probe_accuracy = tensorflow.keras.metrics.BinaryAccuracy(name="p_acc")

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