如何解决BiLSTM-Masked CRF模型编译时的损失函数调用错误？

问题背景 #

基于TensorFlow Keras构建BiLSTM-Masked CRF模型，Masked CRF采用自定义实现。完成模型结构搭建后，编译阶段触发如下TypeError：

TypeError                                 Traceback (most recent call last)
<ipython-input-185-89009da12e28> in <cell line: 1>()
----> 1 loss,_,_ = masked_crf.decode(logits=dense_layer, label_ids=y, lengths=MAX_LEN)
2 base_model.compile(
3     optimizer="adam",
4     loss=loss
5 )

5 frames
/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py in error_handler(*args, **kwargs)
68             # To get the full stack trace, call:
69             # `tf.debugging.disable_traceback_filtering()`
---> 70             raise e.with_traceback(filtered_tb) from None
71         finally:
72             del filtered_tb

TypeError: Exception encountered when calling layer "tf.cond_19" (type TFOpLambda).

To be compatible with tf.function, Python functions must return zero or more Tensors or ExtensionTypes or None values; in compilation of <function crf_sequence_score.<locals>._single_seq_fn at 0x7a396c1249d0>, found return value of type KerasTensor, which is not a Tensor or ExtensionType.

Call arguments received by layer "tf.cond_19" (type TFOpLambda):
• pred=tf.Tensor(shape=(), dtype=bool)
• true_fn=<function crf_sequence_score.<locals>._single_seq_fn at 0x7a396c1249d0>
• false_fn=<function crf_sequence_score.<locals>._multi_seq_fn at 0x7a396dac11b0>
• name=None

错误根源 #

1. Keras Layer实现不规范：自定义MaskedCRF层未实现标准的call方法，而是用自定义decode方法处理逻辑，导致KerasTensor无法被tf.function正确解析。
2. 可训练参数定义位置错误：CRF转移矩阵trans在decode方法内定义，不属于Layer的可训练参数集合，引发张量类型不兼容。
3. 损失函数使用错误：直接将decode返回的损失张量传入compile，不符合Keras要求的“接收y_true和y_pred的损失函数”规范。

修复方案 #

1. 修改MaskedCRF层实现 #

将CRF逻辑整合到标准call方法中，区分训练/推理模式，将转移矩阵作为Layer的可训练参数：

import numpy as np
import tensorflow as tf
import tensorflow_addons as tfa

class MaskedCRF(tf.keras.layers.Layer):
    def __init__(self, num_output, use_mask, label2idx_map, **kwargs):
        super(MaskedCRF, self).__init__(**kwargs)
        self.num_output = num_output
        self.label2idx_map = label2idx_map
        self.use_mask = use_mask

        # 初始化CRF转移矩阵（作为可训练参数）
        self.trans = tf.Variable(
            initial_value=tf.random.normal(shape=[self.num_output, self.num_output], stddev=0.02),
            trainable=True,
            name="crf_transitions"
        )

        self.mask_tran_matrix = None
        if use_mask:
            self.mask_tran_matrix = self.get_mask_trans()

    def call(self, inputs, training=None, label_ids=None, lengths=None):
        logits = inputs
        # 应用转移矩阵掩码
        if self.mask_tran_matrix is not None:
            trans = tf.minimum(self.trans, self.mask_tran_matrix)
        else:
            trans = self.trans

        if training and label_ids is not None:
            # 训练模式：计算损失
            log_likelihood, _ = tfa.text.crf_log_likelihood(
                inputs=logits,
                tag_indices=label_ids,
                sequence_lengths=lengths,
                transition_params=trans
            )
            per_example_loss = -log_likelihood
            loss = tf.math.reduce_mean(per_example_loss)
            # 将损失添加到模型
            self.add_loss(loss)
            # 返回预测结果（可选，用于训练过程中的评估）
            label_pred, _, _ = tfa.text.crf_decode(potentials=logits, transition_params=trans, sequence_length=lengths)
            return label_pred
        else:
            # 推理模式：直接返回预测结果
            label_pred, _, _ = tfa.text.crf_decode(potentials=logits, transition_params=trans, sequence_length=lengths)
            return label_pred

    def get_mask_trans(self):
        size = len(self.label2idx_map)
        tag_lst = self.label2idx_map.keys()

        mask_mat = np.ones(shape=(size, size), dtype=np.float32)
        mask_tran_matrix = np.ones(shape=(size, size), dtype=np.float32)

        is_scheme_bioes = False
        flag_e = False
        flag_s = False
        for tag in tag_lst:
            if tag.startswith("E-"):
                flag_e = True
            if tag.startswith("S-"):
                flag_s = True

        if flag_e and flag_s:
            is_scheme_bioes = True
            print("检测到BIOES标注格式。")
        else:
            print("检测到BIO标注格式。")

        for col_tag, col_index in self.label2idx_map.items():
            if col_tag.startswith("I-"):
                slot_name = col_tag.replace("I-", "")
                begin_slot = "B-" + slot_name
                for row_tag, row_index in self.label2idx_map.items():
                    if is_scheme_bioes:
                        if row_tag != begin_slot and row_tag != col_tag:
                            row_index = min(row_index, size - 1)
                            col_index = min(col_index, size - 1)
                            mask_tran_matrix[row_index, col_index] = -1.0
                    else:
                        if row_tag != col_tag:
                            row_index = min(row_index, size - 1)
                            col_index = min(col_index, size - 1)
                            mask_tran_matrix[row_index, col_index] = -1.0

        mask_mat = tf.convert_to_tensor(mask_mat, dtype=tf.float32)
        mask_tran_matrix = tf.convert_to_tensor(mask_tran_matrix, dtype=tf.float32)
        result_mask = tf.minimum(mask_mat, mask_tran_matrix)

        print("mask_mat shape:", mask_mat.shape)
        print("mask_tran_matrix shape:", mask_tran_matrix.shape)

        return 100 * result_mask

2. 修改模型构建与编译逻辑 #

使用add_loss传递CRF损失，编译时无需指定loss参数：

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Embedding, LSTM, Dense, Bidirectional

# 假设MAX_LEN、number_words、DIM_EMBEDDINGS、number_tags、tag2idx已定义
input_layer = Input(shape=(MAX_LEN,), dtype='int32')
# 标签输入（训练时需要）
label_input = Input(shape=(MAX_LEN,), dtype='int32')
# 序列长度输入（如果需要动态长度，否则可以固定为MAX_LEN）
length_input = Input(shape=(), dtype='int32')

embedding_layer = Embedding(
    input_dim=number_words+1,
    input_length=MAX_LEN,
    output_dim=DIM_EMBEDDINGS,
    trainable=True
)(input_layer)

bilstm_layer = Bidirectional(LSTM(
    units=DIM_EMBEDDINGS,
    return_sequences=True,
    dropout=0.5,
    recurrent_dropout=0.5
))(embedding_layer)

lstm_layer = LSTM(
    units=DIM_EMBEDDINGS*2,
    return_sequences=True,
    dropout=0.5,
    recurrent_dropout=0.5
)(bilstm_layer)

dense_layer = Dense(number_tags, activation='softmax')(lstm_layer)

masked_crf = MaskedCRF(num_output=number_tags, use_mask=True, label2idx_map=tag2idx)
# 训练时传入label_input和length_input
output = masked_crf(dense_layer, training=True, label_ids=label_input, lengths=length_input)

# 模型包含输入层、标签层、长度层
base_model = Model(inputs=[input_layer, label_input, length_input], outputs=output)
base_model.compile(optimizer="adam")

base_model.summary()

3. 训练时的数据输入 #

训练时需要传入三个输入：文本序列、标签序列、序列长度：

# 假设x_train是文本序列，y_train是标签序列，lengths_train是每个序列的实际长度
base_model.fit(
    x=[x_train, y_train, lengths_train],
    y=None,  # 因为损失已经通过add_loss添加
    epochs=10,
    batch_size=32
)

关键说明 #

• 自定义Layer必须实现call方法，这是Keras层的标准接口，确保tf.function能正确处理张量类型。
• 可训练参数必须在__init__中定义，才能被Keras的参数跟踪机制识别。
• 使用add_loss方法添加自定义损失，避免直接传入损失张量到compile的错误用法。

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