TensorFlow自定义带class-fam映射的损失函数报错修复求助

问题分析 #

报错的核心原因是：在map_class_to_fam函数中，直接将TensorFlow张量作为Python字典的键来查找家族类。但TensorFlow的张量是不可哈希对象，无法直接用作字典的键，这就导致了TypeError: Tensor is unhashable的错误；即使切换到V1模式，张量也不能被Python字典正确解析，进而触发KeyError。

修复方案 #

将类到家族类的映射关系转换为TensorFlow的常量张量，使用TensorFlow原生的tf.gather操作完成映射，全程保持计算图兼容，避免Python字典与张量的冲突。

完整修复代码 #

import tensorflow as tf
import numpy as np
import random

# Generate synthetic dataset
num_samples     = 1000
num_classes     = 50
num_fam_classes = 3

# Create random data
X_train       = np.random.rand(num_samples, 10)
y_train_class = np.random.randint(0, num_classes, size=num_samples)

X_val       = np.random.rand(num_samples, 10)
y_val_class = np.random.randint(0, num_classes, size=num_samples)

# convert to category
y_train_one_hot = tf.keras.utils.to_categorical(y_train_class, num_classes=50)
y_val_one_hot   = tf.keras.utils.to_categorical(y_val_class, num_classes=50)

# map between class and fam class - 转换为TensorFlow常量张量
class_to_fam_array = np.array([random.randint(0, 2) for _ in range(num_classes)])
class_to_fam_tensor = tf.constant(class_to_fam_array, dtype=tf.int64)

# define loss:
def custom_loss(y_true, y_pred):
    # --- step 1 : 类别交叉熵损失
    ce_loss_class = tf.keras.losses.categorical_crossentropy(y_true, y_pred, from_logits=False)

    # --- step 2 : 家族类惩罚项计算
    # 获取真实类别和预测类别索引
    true_classes = tf.argmax(y_true, axis=1, output_type=tf.int64)
    predicted_classes = tf.argmax(y_pred, axis=1, output_type=tf.int64)

    # 通过tf.gather完成类别到家族类的映射
    true_fam_classes = tf.gather(class_to_fam_tensor, true_classes)
    predicted_fam_classes = tf.gather(class_to_fam_tensor, predicted_classes)

    # 计算家族类不一致时的惩罚
    # 当家族类不同时，惩罚项取固定值（可根据需求调整权重）
    fam_mismatch = tf.cast(tf.not_equal(true_fam_classes, predicted_fam_classes), dtype=tf.float32)
    penalty = fam_mismatch * 1.0

    # 总损失 = 类别损失 + 惩罚项
    total_loss = ce_loss_class + penalty

    return total_loss

# model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(50, activation='softmax')
])

model.compile(optimizer='adam', loss=custom_loss, metrics=['accuracy'])
model.fit(X_train, y_train_one_hot, epochs=10, batch_size=32, validation_data=[X_val, y_val_one_hot])

关键修改说明 #

• 映射方式替换：把Python字典class_to_fam_dict换成TensorFlow常量张量class_to_fam_tensor，用tf.gather实现批量的类别到家族类的映射，完全兼容计算图模式。
• 惩罚项简化：直接判断家族类是否匹配，用固定权重作为惩罚（可根据需求调整数值），比计算家族类的交叉熵更高效，也能满足“家族类错误时加重惩罚”的需求。
• 类型统一：指定tf.argmax的output_type=tf.int64，确保与class_to_fam_tensor的 dtype 一致，避免类型不匹配问题。

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