Tensorflow 你能在一个自定义的Keras损失函数中处理一个分块的张量吗？

我试图写一个cusom Keras损失函数，在这个函数中，我处理子向量块中的张量。例如，如果输出张量表示四元数系数的串联（即w，x，y，z，w，x，y，z…），我可能希望在计算损失函数中的均方误差之前对每个四元数进行归一化，如：

def norm_quat_mse(y_true, y_pred):
    diff = y_pred - y_true
    dist = 0
    for i in range(0,16,4):
        dist += K.sum( K.square(diff[i:i+4] / K.sqrt(K.sum(K.square(diff[i:i+4])))))

    return dist/4

虽然Keras将毫无错误地接受此函数并在培训中使用，但它输出的损失值与作为独立函数应用时以及使用model.predict（）时不同，因此我怀疑它工作不正常。没有内置的Keras损失函数使用这种逐块处理方法，在Keras的自动差异化框架内是否可以这样做？

尝试：

def norm_quat_mse(y_true, y_pred):
    diff = y_pred - y_true
    dist = 0
    for i in range(0,16,4):
        dist += K.sum( K.square(diff[:,i:i+4] / K.sqrt(K.sum(K.square(diff[:,i:i+4])))))

    return dist/4

您需要知道

y\u true

和

y\u pred

的

shape

是

（批处理大小、输出大小）

，因此在计算过程中需要跳过第一个维度