自定义ArcFace损失函数TensorFlow数据类型错误求助

问题根源 #

你的代码混用了PyTorch与TensorFlow的API（如size()、scatter_、直接索引赋值），同时部分张量操作不符合TensorFlow的规范，导致数据类型不匹配错误。

修复后的完整代码 #

import tensorflow as tf
from tensorflow import Tensor

def call(self, y_true: Tensor, y_pred: Tensor) -> Tensor:
    # 归一化分类器权重与模型输出特征（若你的逻辑中两者角色相反，请调换参数）
    projector = tf.math.l2_normalize(y_true, axis=1)
    predictor = tf.math.l2_normalize(y_pred, axis=1)
    
    # 计算特征与权重的余弦相似度，修正维度匹配
    cos_theta = tf.matmul(predictor, tf.transpose(projector))
    cos_theta = tf.clip_by_value(cos_theta, -1.0, 1.0)
    
    # 筛选有效类别索引（过滤标签为-1的样本）
    valid_mask = tf.where(self.num_classes != -1)
    index = tf.squeeze(valid_mask, axis=1)
    
    # 构造带margin的one-hot矩阵（TensorFlow原生实现）
    batch_size = tf.shape(index)[0]
    num_classes = tf.shape(cos_theta)[1]
    # 生成scatter所需的索引对：[行位置, 类别ID]
    scatter_indices = tf.stack([tf.range(batch_size), tf.gather(self.num_classes, index)], axis=1)
    m_hot = tf.scatter_nd(scatter_indices, tf.fill([batch_size], self.margin), [batch_size, num_classes])
    
    # 对目标类别添加角度margin
    theta = tf.acos(cos_theta)
    theta = tf.tensor_scatter_nd_add(theta, valid_mask, tf.gather(m_hot, tf.range(batch_size)))
    
    # 转换回余弦值并应用缩放因子
    cos_theta_margin = tf.cos(theta)
    cos_theta_scaled = tf.multiply(cos_theta_margin, self.scale)
    
    # 计算标准ArcFace交叉熵损失
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=tf.gather(self.num_classes, index), 
        logits=cos_theta_scaled
    )
    return tf.reduce_mean(loss)

核心修复细节 #

• API统一：将PyTorch风格的size()替换为TensorFlow的tf.shape()，scatter_替换为tf.scatter_nd，直接索引赋值替换为tf.tensor_scatter_nd_add（适配TensorFlow张量不可变特性）。
• 类型对齐：所有张量操作自动继承输入 dtype，避免手动指定类型导致的不匹配；索引操作使用TensorFlow原生支持的整数类型。
• 维度修正：调整矩阵乘法顺序，确保特征与权重的维度正确匹配，避免计算错误。
• 损失逻辑修正：ArcFace损失需基于交叉熵计算，原代码直接返回softmax结果不符合损失函数定义。

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